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:
- Deploy an Autopilot cluster.
- Create a custom ComputeClass which supports the accelerator type (TPU or GPU)
- Create a ResourceClaimTemplate for GPUs (RDMA) or non-GPU (TPU)
- 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:
- code_block
- <ListValue: [StructValue([(‘code’, ‘export PROJECT_ID=$(gcloud config get project) #automatically sets your Project_IDrnexport REGION=”REGION”rnexport CLUSTER_NAME=”CLUSTER_NAME”rnexport NETWORK=”NETWORK”rnexport SUBNETWORK=”SUBNETWORK”rnexport RESERVATION_URL=”RESERVATION_URL”rnexport HF_TOKEN=”HUGGING_FACE_TOKEN”‘), (‘language’, ”), (‘caption’, <wagtail.rich_text.RichText object at 0x7fb066260250>)])]>
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.
1. Deploy an Autopilot cluster
Deploy an Autopilot cluster.
- code_block
- <ListValue: [StructValue([(‘code’, ‘gcloud container clusters create-auto $CLUSTER_NAME \rn –project=$PROJECT_ID \rn –region=$REGION \rn –release-channel=rapid \rn –network=$NETWORK \rn –subnetwork=$SUBNETWORK’), (‘language’, ”), (‘caption’, <wagtail.rich_text.RichText object at 0x7fb066260970>)])]>
2. Create a custom ComputeClass
Example: GPU B200 custom ComputeClass with managed DRANET support and a reservation.
- code_block
- <ListValue: [StructValue([(‘code’, ‘apiVersion: cloud.google.com/v1rnkind: ComputeClassrnmetadata:rn name: dranet-a4-computeclassrnspec:rn nodePoolAutoCreation:rn enabled: truern nodePoolConfig:rn dra:rn networking:rn enabled: truern priorities:rn – machineType: a4-highgpu-8grn gpu:rn count: 8rn type: nvidia-b200rn acceleratorNetworkProfile: autorn reservations:rn affinity: Specificrn specific:rn – name: ${RESERVATION_URL}rn project: ${PROJECT_ID}’), (‘language’, ”), (‘caption’, <wagtail.rich_text.RichText object at 0x7fb066260b50>)])]>
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.
- code_block
- <ListValue: [StructValue([(‘code’, ‘apiVersion: cloud.google.com/v1rnkind: ComputeClassrnmetadata:rn name: dra-gke-autornspec:rn nodePoolAutoCreation:rn enabled: truern nodePoolConfig:rn dra:rn networking:rn enabled: truern priorities:rn – tpu:rn type: tpu-v6e-slicern count: 8rn topology: “2×4″ rn acceleratorNetworkProfile: autorn location:rn zones: rn – us-east5-b’), (‘language’, ”), (‘caption’, <wagtail.rich_text.RichText object at 0x7fb066260430>)])]>
3. Create a ResourceClaimTemplate
RDMA support deviceClassName: mrdma.google.com ResourceClaimTemplate example for GPUs:
- code_block
- <ListValue: [StructValue([(‘code’, ‘apiVersion: resource.k8s.io/v1rnkind: ResourceClaimTemplaternmetadata:rn name: all-mrdmarnspec:rn spec:rn devices:rn requests:rn – name: req-mrdmarn exactly:rn deviceClassName: mrdma.google.comrn allocationMode: All’), (‘language’, ”), (‘caption’, <wagtail.rich_text.RichText object at 0x7fb0662602b0>)])]>
Non-RDMA deviceClassName: netdev.google.com ResourceClaimTemplate example for TPUs.
- code_block
- <ListValue: [StructValue([(‘code’, ‘apiVersion: resource.k8s.io/v1rnkind: ResourceClaimTemplaternmetadata:rn name: all-netdevrnspec:rn spec:rn devices:rn requests:rn – name: req-netdevrn exactly:rn deviceClassName: netdev.google.comrn allocationMode: All’), (‘language’, ”), (‘caption’, <wagtail.rich_text.RichText object at 0x7fb066260310>)])]>
4. Deploy workload and reference ComputeClass and ResourceClaim
Create a secret in your cluster
- code_block
- <ListValue: [StructValue([(‘code’, ‘kubectl create secret generic hf-secret \rn –from-literal=hf_token=${HF_TOKEN}’), (‘language’, ”), (‘caption’, <wagtail.rich_text.RichText object at 0x7fb0662601f0>)])]>
Example deploying GPUs
- code_block
- <ListValue: [StructValue([(‘code’, ‘apiVersion: apps/v1rnkind: Deploymentrnmetadata:rn name: gemma-4-31-deployrnspec:rn replicas: 2rn selector:rn matchLabels:rn app: gemma4rn template:rn metadata:rn labels:rn app: gemma4rn ai.gke.io/model: gemma-4-31brn ai.gke.io/inference-server: vllmrn spec:rn resourceClaims:rn – name: rdma-claim rn resourceClaimTemplateName: all-mrdmarn containers:rn – name: vllm-inferencern image: us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:gemma4rn resources:rn requests:rn cpu: “10”rn memory: “1000Gi”rn ephemeral-storage: “1Ti”rn nvidia.com/gpu: “8”rn limits:rn cpu: “10”rn memory: “1000Gi”rn ephemeral-storage: “1Ti”rn nvidia.com/gpu: “8”rn claims:rn – name: rdma-claimrn command: [“python3”, “-m”, “vllm.entrypoints.openai.api_server”]rn args:rn – –model=$(MODEL_ID)rn – –tensor-parallel-size=8rn – –host=0.0.0.0rn – –port=8000rn – –max-model-len=131072rn – –max-num-seqs=16rn – –enable-chunked-prefillrn – –gpu-memory-utilization=0.90rn env:rn – name: MODEL_IDrn value: google/gemma-4-31Brn – name: HUGGING_FACE_HUB_TOKENrn valueFrom:rn secretKeyRef:rn name: hf-secretrn key: hf_tokenrn volumeMounts:rn – mountPath: /dev/shmrn name: dshmrn startupProbe:rn httpGet:rn path: /healthrn port: 8000rn failureThreshold: 240rn periodSeconds: 10rn livenessProbe:rn httpGet:rn path: /healthrn port: 8000rn periodSeconds: 10rn readinessProbe:rn httpGet:rn path: /healthrn port: 8000rn periodSeconds: 5rn volumes:rn – name: dshmrn emptyDir:rn medium: Memoryrn nodeSelector:rn cloud.google.com/compute-class: dranet-a4-computeclass’), (‘language’, ”), (‘caption’, <wagtail.rich_text.RichText object at 0x7fb0662601c0>)])]>
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.
Next Steps
Take a deeper dive into GKE managed DRANET and autopilot with these resources:
- Hands-on Lab: GKE Autopilot clusters with TPUs, GKE managed DRANET and Gemma 4
- Document set: DRANET
- Documentation: AI Hypercomputer
Want to ask a question, find out more, or share a thought? Please connect with me on Linkedin.
