gke-a3-highgpu.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-highgpu

vars:
  # The following variables should be over-written in the deployment.yaml file.
  # Your GCP Project ID
  project_id: ## Set GCP Project ID Here ##

  # This should be unique across all of your Cluster
  # Toolkit Deployments.
  deployment_name: gke-a3-high

  # The GCP Region used for this deployment.
  region: us-central1

  # The GCP Zone used for this deployment.
  zone: us-central1-c

  # Cidr block containing the IP of the machine calling terraform.
  # The following line must be updated for this example to work.
  authorized_cidr: <your-ip-address>/32

  gcp_public_cidrs_access_enabled: false
  kueue_configuration_path: $(ghpc_stage("./kueue-configuration.yaml.tftpl"))
  nccl_tcpx_installer_path: $(ghpc_stage("./nccl-tcpx-installer.yaml.tftpl"))
  nccl_config_path: $(ghpc_stage("./nccl-config.yaml"))
  nri_device_injector_path: $(ghpc_stage("./nri-device-injector.yaml"))

  network_name: $(vars.deployment_name)-net

  # The number of nodes to be created
  static_node_count: 2

  # 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

  accelerator_type: nvidia-h100-80gb
  version_prefix: "1.35."
  nccl_tcpx_version: v3.1.9

deployment_groups:
- group: primary
  modules:
  - id: network1
    source: modules/network/vpc
    settings:
      network_name: $(vars.network_name)
      ips_per_nat: 6
      mtu: 8244
      subnetworks:
      - subnet_name: $(vars.deployment_name)-subnet
        subnet_region: $(vars.region)
        subnet_ip: 192.168.0.0/18
      secondary_ranges_list:
      - subnetwork_name: $(vars.deployment_name)-subnet
        ranges:
        - range_name: pods
          ip_cidr_range: 10.4.0.0/14
        - range_name: services
          ip_cidr_range: 10.0.32.0/20

  - 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: gpunets
    source: modules/network/multivpc
    settings:
      network_name_prefix: $(vars.deployment_name)-gpunet
      global_ip_address_range: 192.169.0.0/16
      network_count: 4
      subnetwork_cidr_suffix: 20 # the subnet can support upto 4k nodes
      mtu: 8244

  - id: gke_cluster
    source: modules/scheduler/gke-cluster
    use: [network1, gpunets, workload_service_account]
    settings:
      enable_private_endpoint: false  # Allows for access from authorized public IPs
      gcp_public_cidrs_access_enabled: $(vars.gcp_public_cidrs_access_enabled)
      configure_workload_identity_sa: true
      master_authorized_networks:
      - cidr_block: $(vars.authorized_cidr)  # Allows your machine run kubectl command. It's required for the multi-network setup.
        display_name: "kubectl-access-network"
      k8s_network_names:
        gvnic_prefix: vpc
        gvnic_start_index: 1
      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]

  - id: a3_highgpu_pool
    source: modules/compute/gke-node-pool
    use: [gke_cluster, gpunets, node_pool_service_account]
    settings:
      install_gpu_direct_manifests: false
      run_workload_script: false
      machine_type: a3-highgpu-8g
      auto_upgrade: true
      static_node_count: $(vars.static_node_count)
      zones: [$(vars.zone)]
      guest_accelerator:
      - type: $(vars.accelerator_type)
        count: 8
      reservation_affinity:
        consume_reservation_type: SPECIFIC_RESERVATION
        specific_reservations:
        - name: $(vars.reservation)
    outputs: [instructions]

  - id: workload_component_install
    source: modules/management/kubectl-apply
    use: [gke_cluster]
    settings:
      apply_manifests:
      - source: $(vars.nccl_config_path)
      - source: $(vars.nri_device_injector_path)
      - source: $(vars.nccl_tcpx_installer_path)
        template_vars:
          accelerator_type: $(vars.accelerator_type)
          version: $(vars.nccl_tcpx_version)
      kueue:
        install: true
        wait: true
        version: "0.17.1"
        config_path: $(vars.kueue_configuration_path)
        config_template_vars:
          num_gpus: $(a3_highgpu_pool.static_gpu_count)
          accelerator_type: $(vars.accelerator_type)
      jobset:
        install: true