Create an A4X Max GKE cluster

This document explains how to deploy a Google Kubernetes Engine (GKE) cluster that uses A4X Max instances. This setup is ideal for running workloads for machine learning (ML) training and high performance computing (HPC).

The A4X Max machine series runs on an exascale platform based on NVIDIA's rack-scale architecture and is optimized for compute and memory-intensive, network-bound ML training and HPC workloads. A4X Max offers bare metal instances, which provide direct access to the host server's CPU and memory, without the Compute Engine hypervisor layer.

To learn more about the A4X machine series, see the A4X Max and A4X machine series section in the Compute Engine documentation.

Before you begin

Before you start, make sure that you have performed the following tasks:

Enable the Google Kubernetes Engine API.

Enable Google Kubernetes Engine API

To use the Google Cloud CLI for this task, install and then initialize the gcloud CLI. If you previously installed the gcloud CLI, get the latest version by running the gcloud components update command. Earlier gcloud CLI versions might not support running the commands in this document.
Note: For existing gcloud CLI installations, make sure to set the compute/region property. If you use primarily zonal clusters, set the compute/zone instead. By setting a default location, you can avoid errors in the gcloud CLI like the following: One of [--zone, --region] must be supplied: Please specify location. You might need to specify the location in certain commands if the location of your cluster differs from the default that you set.

Required roles

To get the permissions that you need to deploy the cluster, ask your administrator to grant you the following IAM roles:

Kubernetes Engine Admin (roles/container.admin) on the project
Compute Admin (roles/compute.admin) on the project
Storage Admin (roles/storage.admin) on the project
Project IAM Admin (roles/resourcemanager.projectIamAdmin) on the project
Service Account Admin (roles/iam.serviceAccountAdmin) on the project
Service Account User (roles/iam.serviceAccountUser) on the project
Service Usage Consumer (roles/serviceusage.serviceUsageConsumer) on the project
To consume a single-project reservation: Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1) on the project
To consume a shared reservation: Compute Instance Admin (v1) (roles/compute.instanceAdmin.v1) on both the owner project and any consumer project where you want to consume the reservation

For more information about granting roles, see Manage access to projects, folders, and organizations.

You might also be able to get the required permissions through custom roles or other predefined roles.

Choose a consumption option and obtain capacity

To obtain capacity, complete the following steps:

Choose a consumption option: Make your choice based on how you want to get and use GPU resources. To learn more, see Choose a consumption option.

For GKE, consider the following additional information:
- A4X Max compute instances cannot be provisioned by using flex-start.
- You must use the reservation-bound provisioning model to create clusters with A4X Max. Other provisioning models are not supported.
Obtain capacity: The process to obtain capacity differs for each consumption option. To learn about the process for your chosen consumption option, see Capacity overview.

Note: When you request A4X Max capacity, you obtain it in all capacity mode. This mode is the only supported reservation operational mode for A4X Max machine types. For more information about all capacity mode, see Reservation operational mode.

Requirements

The following requirements apply to an AI-optimized GKE cluster that uses A4X Max instances:

Verify that you use GKE version 1.35.0-gke.2745000 or later (for 1.35), or version 1.34.3-gke.1318000 or later (for 1.34). These versions help ensure that A4X Max uses the following:
- R580.95.05, which is the minimum GPU driver version for the GB300 GPUs in A4X Max bare metal instances. This driver is enabled by default.
- Coherent Driver-based Memory Management (CDMM), which is enabled by default. NVIDIA recommends that Kubernetes clusters enable this mode to resolve memory over-reporting. CDMM lets you manage GPU memory through the driver instead of the operating system (OS). This approach helps you avoid OS onlining of GPU memory, and exposes the GPU memory as a Non-Uniform Memory Access (NUMA) node to the OS. Multi-instance GPUs are not supported when CDMM is enabled. For more information about CDMM, see Hardware and Software Support.
- GPUDirect RDMA, which is recommended to let A4X Max node pools use the networking capabilities of A4X Max.
To use GPUDirect RDMA, the GKE nodes must use a Container-Optimized OS node image. Ubuntu and Windows node images are not supported.

Customize the DRANET configuration

This GKE blueprint for A4X Max automatically configures GKE-managed Dynamic Resource Allocation for Networking (DRANET).

When you provision a node pool with DRANET enabled, Cluster Toolkit automatically detects the machine type and requests the Multi-Host Remote Direct Memory Access (MRDMA) interfaces (mrdma.google.com). Cluster Toolkit also creates a ResourceClaimTemplate object named default-rdma in your cluster. If you configure multiple node pools with DRANET, then Cluster Toolkit automatically manages the deployment of templates to prevent conflicts between duplicate resources.

To allocate network devices or change device classes, you can customize the template by using the dranet_allocation_mode, dranet_device_count, and dranet_device_class_name variables in your node pool module settings.

The following example shows how to configure DRANET in a node pool:

  - id: a4x-max-pool
    source: modules/compute/gke-node-pool
    settings:
      enable_dranet: true
      # Optional: Override DRANET template defaults
      # dranet_allocation_mode: ExactCount
      # dranet_device_count: 4

To allocate DRANET resources, do one of the following depending on your deployment method:

If you use Cluster Toolkit workload templates, such as the gke-job-template module, then you don't need to configure the resources manually. The toolkit automatically injects the default resource claims into your workloads.
If you deploy a custom workload by using a Kubernetes YAML manifest, then you must explicitly configure the claims in your manifest. To configure the claims, do the following:
- Declare the claim at the Pod level: In the spec.template.spec.resourceClaims section, add a claim that references the default-rdma template:
```
spec:
  template:
    spec:
      resourceClaims:
      - name: rdma
        resourceClaimTemplateName: default-rdma
```
- Bind the claim at the container level: In the resources.claims section in your container definition, add a claim to bind the Pod-level claim that you declared:
```
resources:
  claims:
  - name: rdma
```

To view examples of custom workload manifests, see Deploy a workload with RDMA resources or the gke-a4x-max-bm/nccl-jobset-example.yaml example on GitHub.