gke-a3-ultragpu.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: gke-a3-ultra

vars:
  project_id: # add this
  deployment_name: # add this
  region: # add this
  zone: # add this
  # Cidr block containing the IP of the machine calling terraform.
  # The following line must be updated for this example to work.
  authorized_cidr: # add this
  # The name of the compute engine reservation in the form of
  # <reservation-name>
  # To target a BLOCK_NAME, the name of the extended reservation
  # can be inputted as <reservation-name>/reservationBlocks/<reservation-block-name>
  reservation: # add this
  static_node_count: # add this

  kueue_configuration_path: $(ghpc_stage("./kueue-configuration.yaml.tftpl"))
  gib_installer_path: $(ghpc_stage("./nccl-installer.yaml.tftpl"))
  system_node_pool_disk_size_gb: 200
  a3ultra_node_pool_disk_size_gb: 100
  accelerator_type: nvidia-h200-141gb
  version_prefix: "1.35."

  # # To enable Managed-Lustre please uncomment this section and fill out the settings.
  # # Additionally, please uncomment the private_service_access, lustre_firewall_rule, managed-lustre and lustre-pv modules.
  # # 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: gke-lustre-instance

  # # The values of size_gib and per_unit_storage_throughput are co-related
  # # Please refer https://cloud.google.com/managed-lustre/docs/create-instance#performance-tiers
  # # Storage capacity of the lustre instance in GiB
  # lustre_size_gib: 36000

  # # Maximum throughput of the lustre instance in MBps per TiB
  # per_unit_storage_throughput: 500

deployment_groups:
- group: primary
  modules:
  - id: gke-a3-ultra-net-0
    source: modules/network/vpc
    settings:
      network_name: $(vars.deployment_name)-net-0
      ips_per_nat: 6
      subnetworks:
      - subnet_name: $(vars.deployment_name)-sub-0
        subnet_region: $(vars.region)
        subnet_ip: 192.168.0.0/18
      secondary_ranges_list:
      - subnetwork_name: $(vars.deployment_name)-sub-0
        ranges:
        - range_name: pods
          ip_cidr_range: 10.4.0.0/14
        - range_name: services
          ip_cidr_range: 10.0.32.0/20
      firewall_rules:
      - name: $(vars.deployment_name)-internal-0
        ranges: [192.168.0.0/16]
        allow:
        - protocol: tcp
          ports: ["0-65535"]
        - protocol: udp
          ports: ["0-65535"]
        - protocol: icmp

  - id: gke-a3-ultra-net-1
    source: modules/network/vpc
    settings:
      network_name: $(vars.deployment_name)-net-1
      ips_per_nat: 6
      subnetworks:
      - subnet_name: $(vars.deployment_name)-sub-1
        subnet_region: $(vars.region)
        subnet_ip: 192.168.64.0/18
      firewall_rules:
      - name: $(vars.deployment_name)-internal-1
        ranges: [192.168.0.0/16]
        allow:
        - protocol: tcp
          ports: ["0-65535"]
        - protocol: udp
          ports: ["0-65535"]
        - protocol: icmp

  - id: gke-a3-ultra-rdma-net
    source: modules/network/gpu-rdma-vpc
    settings:
      network_name: $(vars.deployment_name)-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: $(vars.deployment_name)-rdma-sub
        count: 8
        ip_range: 192.168.128.0/18
        region: $(vars.region)

  - id: node_pool_service_account
    source: modules/project/service-account
    settings:
      name: gke-np-sa
      project_roles:
      - logging.logWriter
      - monitoring.metricWriter
      - monitoring.viewer
      - stackdriver.resourceMetadata.writer
      - storage.objectViewer
      - artifactregistry.reader

  - id: workload_service_account
    source: modules/project/service-account
    settings:
      name: gke-wl-sa
      project_roles:
      - logging.logWriter
      - monitoring.metricWriter
      - monitoring.viewer
      - stackdriver.resourceMetadata.writer
      - storage.objectAdmin
      - artifactregistry.reader
      - container.admin

  - id: training_bucket
    source: modules/file-system/cloud-storage-bucket
    settings:
      local_mount: /training-data
      name_prefix: training
      random_suffix: true
      force_destroy: false
      enable_hierarchical_namespace: true

  - id: checkpoint_bucket
    source: modules/file-system/cloud-storage-bucket
    settings:
      local_mount: /checkpoint-data
      name_prefix: checkpoint
      random_suffix: true
      force_destroy: false
      enable_hierarchical_namespace: true

  - id: a3-ultragpu-cluster
    source: modules/scheduler/gke-cluster
    use: [gke-a3-ultra-net-0, workload_service_account]
    settings:
      system_node_pool_machine_type: "e2-standard-16"
      system_node_pool_disk_size_gb: $(vars.system_node_pool_disk_size_gb)
      system_node_pool_taints: []
      enable_dcgm_monitoring: true
      enable_gcsfuse_csi: true
      enable_managed_lustre_csi: true # Enable Managed Lustre for the cluster
      enable_private_endpoint: false # Allows access from authorized public IPs
      configure_workload_identity_sa: true
      master_authorized_networks:
      - cidr_block: $(vars.authorized_cidr) # Allows your machine to run the kubectl command. Required for multi network setup.
        display_name: "kubectl-access-network"
      additional_networks:
        $(concat(gke-a3-ultra-net-1.instance_additional_networks,
         gke-a3-ultra-rdma-net.subnetwork_interfaces_gke
        ))
      # Cluster versions cannot be updated through the toolkit after creation
      # Please manage cluster version from the Google Cloud Console directly
      version_prefix: $(vars.version_prefix)
      release_channel: REGULAR
      maintenance_exclusions:
      - name: no-minor-or-node-upgrades-indefinite
        exclusion_scope: NO_MINOR_OR_NODE_UPGRADES
        exclusion_end_time_behavior: UNTIL_END_OF_SUPPORT
    outputs: [instructions]

  # # --- MANAGED LUSTRE ADDITIONS ---
  # # Private Service Access (PSA) requires the compute.networkAdmin role which is
  # # included in the Owner role, but not Editor.
  # # PSA is required for all Managed Lustre functionality.
  # # https://cloud.google.com/vpc/docs/configure-private-services-access#permissions
  # - id: private_service_access
  #   source: modules/network/private-service-access
  #   use: [gke-a3-ultra-net-0]
  #   settings:
  #     prefix_length: 22

  # # Firewall to allow Managed Lustre connection
  # - id: lustre_firewall_rule
  #   source: modules/network/firewall-rules
  #   use: [gke-a3-ultra-net-0]
  #   settings:
  #     ingress_rules:
  #     - name: $(vars.deployment_name)-allow-lustre-traffic
  #       description: Allow Managed Lustre traffic
  #       source_ranges:
  #       - $(private_service_access.cidr_range)
  #       allow:
  #       - protocol: tcp
  #       ports:
  #       - "988"

  # - id: managed-lustre
  #   source: modules/file-system/managed-lustre
  #   use: [gke-a3-ultra-net-0, private_service_access]
  #   settings:
  #     name: $(vars.lustre_instance_id)
  #     local_mount: /lustre
  #     remote_mount: lustrefs
  #     size_gib: $(vars.lustre_size_gib)
  #     per_unit_storage_throughput: $(vars.per_unit_storage_throughput)

  # - id: lustre-pv
  #   source: modules/file-system/gke-persistent-volume
  #   use: [managed-lustre, a3-ultragpu-cluster]
  #   settings:
  #     capacity_gib: $(vars.lustre_size_gib)

  - id: a3-ultragpu-pool
    source: modules/compute/gke-node-pool
    use: [a3-ultragpu-cluster, node_pool_service_account]
    settings:
      machine_type: a3-ultragpu-8g
      auto_upgrade: true
      zones: [$(vars.zone)]
      disk_size_gb: $(vars.a3ultra_node_pool_disk_size_gb)
      static_node_count: $(vars.static_node_count)
      guest_accelerator:
      - type: $(vars.accelerator_type)
        count: 8
      reservation_affinity:
        consume_reservation_type: SPECIFIC_RESERVATION
        specific_reservations:
        - name: $(vars.reservation)
      additional_networks:
        $(concat(gke-a3-ultra-net-1.instance_additional_networks,
         gke-a3-ultra-rdma-net.subnetwork_interfaces_gke
        ))
    outputs: [instructions]

  - id: workload-manager-install
    source: modules/management/kubectl-apply
    use: [a3-ultragpu-cluster]
    settings:
      kueue:
        install: true
        config_path: $(vars.kueue_configuration_path)
        config_template_vars:
          num_gpus: $(a3-ultragpu-pool.static_gpu_count)
          accelerator_type: $(vars.accelerator_type)
      jobset:
        install: true
      gib:
        install: true  # NCCL gIB plugin via DaemonSet initContainer
        path: $(vars.gib_installer_path)
        template_vars:
          version: v1.1.2
          accelerator_count: 8

  - id: job-template
    source: modules/compute/gke-job-template
    use: [a3-ultragpu-pool]
    settings:
      image: nvidia/cuda:11.0.3-runtime-ubuntu20.04
      command:
      - nvidia-smi
      node_count: 2
      name: run-nvidia-smi
      k8s_service_account_name: workload-identity-k8s-sa
    outputs: [instructions]

  # When utilizing GCS Fuse Storage Profiles, you must ensure the GKE Service Agent
  # possesses the 'gke.gcsfuse.profileUser' custom IAM role assigned out-of-band.
  # Please refer to 'modules/file-system/gke-persistent-volume/README.md' for instructions.
  # Persistent Volume for training data using GCS Fuse Storage Profile
  - id: training-pv
    source: modules/file-system/gke-persistent-volume
    use: [training_bucket, a3-ultragpu-cluster]
    settings:
      capacity_gib: 1000000
      gcsfuse_storage_class_name: gcsfusecsi-training

  # Persistent Volume for checkpoint data using GCS Fuse Storage Profile
  - id: checkpointing-pv
    source: modules/file-system/gke-persistent-volume
    use: [checkpoint_bucket, a3-ultragpu-cluster]
    settings:
      capacity_gib: 1000000
      gcsfuse_storage_class_name: gcsfusecsi-checkpointing

  # This is an example job that will install and run an `fio`
  # benchmark against the training and checkpointing buckets.
  - id: fio-bench-job-template
    source: modules/compute/gke-job-template
    use: [checkpointing-pv, training-pv, a3-ultragpu-pool]
    settings:
      name: gcsfuse-fio-job
      security_context:  # to make sure the job have enough access to install the fio packages
      - key: runAsUser
        value: 0
      - key: runAsGroup
        value: 100
      - key: fsGroup
        value: 100
      # By adding an ephemeral volume, this will ensure that the job adds:
      # nodeSelector:
      #   cloud.google.com/gke-ephemeral-storage-local-ssd: "true"
      # which is the best practice for using local-ssd for ephemeral storage.
      ephemeral_volumes:
      - type: local-ssd
        mount_path: /scratch-data
        size_gb: 1000  # Use 1 out of 12 TB for local scratch

      k8s_service_account_name: workload-identity-k8s-sa
      image: ubuntu:22.04

      command:
      - bash
      - -c
      - |

        set -eux
        export DEBIAN_FRONTEND=noninteractive

        # Install fio
        apt update -y && apt install -y fio

        # Use a tag to create a unique path for tests
        TAG=`date +%s`

        # Verify mountpoints
        df -h
        mountpoint /scratch-data
        mountpoint /checkpoint-data
        mountpoint /training-data

        # Create temporary directory for fio benchmarks
        mkdir -p /{scratch,training,checkpoint}-data/fio-benchmarks-${TAG}

        # The following will take roughly 10 minutes to complete

        # Perform scratch data write performance test
        fio --ioengine=libaio --filesize=10G --ramp_time=2s --runtime=1m \
          --numjobs=32 --create_serialize=0 --direct=1 --verify=0 \
          --randrepeat=0 --group_reporting --directory=/scratch-data/fio-benchmarks-${TAG} \
          --name=scratch --blocksize=100m --iodepth=64 --readwrite=write

        # Perform training data reading performance test
        fio --ioengine=libaio --filesize=1G --ramp_time=2s --runtime=1m \
          --numjobs=32 --create_serialize=0 --direct=1 --verify=0 \
          --randrepeat=0 --group_reporting --directory=/training-data/fio-benchmarks-${TAG} \
          --name=training --blocksize=1m --iodepth=64 --readwrite=randread

        # Perform checkpoint data writing performance test
        fio --ioengine=libaio --filesize=10G --ramp_time=2s --runtime=1m \
          --numjobs=32 --create_serialize=0 --direct=1 --verify=0 \
          --randrepeat=0 --group_reporting --directory=/checkpoint-data/fio-benchmarks-${TAG} \
          --name=checkpoint --blocksize=100m --iodepth=64 --readwrite=write

        # Perform checkpoint data reading performance test
        fio --ioengine=libaio --filesize=10G --ramp_time=2s --runtime=1m \
          --numjobs=32 --create_serialize=0 --direct=1 --verify=0 \
          --randrepeat=0 --group_reporting --directory=/checkpoint-data/fio-benchmarks-${TAG} \
          --name=checkpoint --blocksize=100m --iodepth=64 --readwrite=read

        # Clean up temporary directories for fio benchmarks
        rm -rf /{scratch,training,checkpoint}-data/fio-benchmarks-${TAG}

    outputs: [instructions]