Google Kubernetes Engine (GKE) managed DRANET supports both GPUs and TPUs. There are several configurations to use this implementation, including standard cluster (where you have full control) and autopilot cluster (where Google does the heavy configs for you). I've been exploring the capabilities and in this blog we will explore setting up for autopilot clusters.

Autopilot and managed DRANET

GKE autopilot is a managed version of GKE that handles nodes, scaling, security, and other preconfigured settings. GKE managed DRANET lets you request and allocate networking resources for your Pods, including network interfaces that support TPUs and Remote Direct Memory Access (RDMA).

Setup flow

To deploy your GKE autopilot cluster and enable managed DRANET, you need to create a Virtual Private Cloud (VPC). Let's walk through the setup:

  1. Deploy an Autopilot cluster.
  2. Create a custom ComputeClass which supports the accelerator type (TPU or GPU)
  3. Create a ResourceClaimTemplate for GPUs (RDMA) or non-GPU (TPU)
  4. Deploy workload and reference the ComputeClass and ResourceClaimTemplate to get the correct networking set up.

Now let's explore the configs for both TPU and GPU.

Configure variables:

Replace the following:

  • REGION: The region where you want to create your cluster, such as us-east1. You can only create the cluster in the region where your reservation or resources exists.
  • CLUSTER_NAME: A name for your cluster, such as dranet-cluster.
  • NETWORK: The name of the VPC network.
  • SUBNETWORK: The name of the subnet in the VPC.
  • RESERVATION_URL: The URL of the reservation that you want to use to create your resources.
  • HUGGING_FACE_TOKEN: The Hugging Face access token to download your model.

2. Create a custom ComputeClass

Example: GPU B200 custom ComputeClass with managed DRANET support and a reservation.

Replace the following:

  • ${RESERVATION} : With the URL of the reservation that you want to use to create your resources.
  • ${PROJECT_ID}: With the ID of the project you are using.

Alternatively you can set the variables in your terminal and use the following command to pass the variables at creation envsubst < filename.yaml | kubectl apply -f -

Example: TPU v6e custom ComputeClass using on-demand example.

3. Create a ResourceClaimTemplate

RDMA support deviceClassName: mrdma.google.com ResourceClaimTemplate example for GPUs:

Non-RDMA deviceClassName: netdev.google.com ResourceClaimTemplate example for TPUs.

4. Deploy workload and reference ComputeClass and ResourceClaim

Create a secret in your cluster

Example deploying GPUs

Notice how the deployment references the ResourceClaimTemplate and ComputeClass. When this kicks off, it triggers a scale-up operation. GKE Autopilot reads the ComputeClass to provision the specific node type and to configure managed DRANET networking. Meanwhile, the resource claim acts as the bridge, binding your Pods directly to the accelerators on those nodes. This process works exactly the same for TPUs.

Posted in