Create a GKE cluster with Pathways

This page shows platform administrators and machine learning (ML) engineers how to create a Google Kubernetes Engine (GKE) cluster configured for Pathways-based workloads. You can provision pre-configured clusters automatically by using Cluster Toolkit, or create clusters manually by using the Google Cloud CLI.

Before you begin

Complete the following prerequisites before you proceed:

Install the gcloud CLI.
Install kubectl and the gke-gcloud-auth-plugin.
If you use Cluster Toolkit, set up Cluster Toolkit and update it to the latest version.
Enable the Kubernetes Engine API.
Verify that you have enough TPU quota or a TPU reservation for the TPU slices that you want to create.

Set up your local environment

gcloud auth login

If you use Cluster Toolkit, also generate Application Default Credentials (ADC), which Terraform uses:

gcloud auth application-default login

Create a GKE cluster

In the following example, you create a cluster with TPU v6e 2x4 node pools. You can create a cluster by using Cluster Toolkit or the gcloud CLI.

Cluster Toolkit

Deploy a GKE TPU cluster by following the instructions in Deploy a GKE TPU v6e cluster. For an overview of all TPU deployment options, see Cloud TPU deployments overview.

When configuring your deployment blueprint, do the following:
- Set enable_pathways_for_tpus to true (enabled by default in Cluster Toolkit TPU blueprints). This setting automatically provisions the required cpu-np CPU node pool (n4-standard-64) and Kueue configurations for the Pathways controller. Cluster Toolkit applies the default Pathways Kueue configuration only if you don't specify a custom Kueue configuration file in the kueue_configuration_path variable.
- Set num_slices to the number of TPU slice node pools required for your workload.
- Set tpu_topology to your selected TPU slice topology (for example, 2x4 or 4x4).
- Choose a consumption option. By default, the TPU v6e deployment file consumes a specific TPU reservation, so specify the reservation name in the reservation_affinity variable. To use a different consumption option, see GKE consumption options.
Verify that the cluster nodes have a STATUS of Ready:
```
kubectl get nodes
```
After you create the cluster, you can submit Pathways workloads by using the gcluster job submit command with the --pathways and --pathways-gcs-location flags. For more information, see Run workloads with Pathways.

gcloud

To create a GKE cluster and set it up for running Pathways workloads, follow these steps:

Set the following environment variables:
```
CLUSTER=GKE_CLUSTER_NAME
PROJECT=PROJECT_ID
ZONE=ZONE
REGION=REGION
CLUSTER_VERSION=GKE_CLUSTER_VERSION
PW_CPU_MACHINE_TYPE="n2-standard-64"
NETWORK=NETWORK
SUBNETWORK=SUB_NETWORK
CLUSTER_NODEPOOL_COUNT=3
TPU_MACHINE_TYPE="ct6e-standard-4t"
WORKERS_PER_SLICE=2
TOPOLOGY="2x4"
NUM_CPU_NODES=1
```
Replace the following:
- GKE_CLUSTER_NAME: the name of your GKE cluster.
- PROJECT_ID: your Google Cloud project ID.
- ZONE: the zone where you create the cluster. Choose a zone where TPU v6e is available. For more information, see Cloud TPU regions and zones.
- REGION: Optional: the region where you create resources if you use a regional cluster.
- GKE_CLUSTER_VERSION: Optional: the GKE cluster version. Use 1.32.2-gke.1475000 or later, and choose a version that's in standard support. For more information, see the GKE release schedule. If you don't specify a version, remove the --cluster-version flag from the following gcloud container clusters create command.
- NETWORK: Optional: the name of a Virtual Private Cloud (VPC) network. The network must exist before you create the cluster. If you don't specify a network, remove the --network flag from the following gcloud container clusters create command.
- SUB_NETWORK: Optional: the name of a subnetwork. The subnetwork must exist before you create the cluster. If you don't specify a subnetwork, remove the --subnetwork flag from the following gcloud container clusters create command.
This example sets default values for the remaining variables, which you can modify for your workload:
- PW_CPU_MACHINE_TYPE: the machine type for the default node pool of the cluster and for the CPU node pool that runs the Pathways head (the Pathways resource manager and proxy).
- NUM_CPU_NODES: the number of nodes in the CPU node pool.
- CLUSTER_NODEPOOL_COUNT: the number of TPU slices. Each slice is a separate node pool.
- TPU_MACHINE_TYPE and TOPOLOGY: the TPU machine type and the TPU topology of each slice.
- WORKERS_PER_SLICE: the number of nodes in each slice, which is the number of TPU chips in the topology divided by the number of chips per node. For example, a 2x4 topology has 8 chips and a ct6e-standard-4t node has 4 chips, so each slice has 2 nodes.

Create a GKE cluster by using the gcloud container clusters create command:

gcloud container clusters create ${CLUSTER} \
    --project=${PROJECT} \
    --zone=${ZONE} \
    --cluster-version=${CLUSTER_VERSION} \
    --scopes=storage-full,gke-default,cloud-platform \
    --machine-type=${PW_CPU_MACHINE_TYPE} \
    --network=${NETWORK} \
    --subnetwork=${SUBNETWORK}

Create the TPU slice node pools by using the gcloud container node-pools create command:

for i in $(seq 1 ${CLUSTER_NODEPOOL_COUNT}); do
    gcloud container node-pools create "tpu-np-${i}" \
        --project=${PROJECT} \
        --zone=${ZONE} \
        --cluster=${CLUSTER} \
        --machine-type=${TPU_MACHINE_TYPE} \
        --num-nodes=${WORKERS_PER_SLICE} \
        --placement-type=COMPACT \
        --tpu-topology=${TOPOLOGY} \
        --scopes=storage-full,gke-default,cloud-platform \
        --workload-metadata=GCE_METADATA
done

These commands don't target a specific reservation. To consume a TPU reservation, add the --reservation-affinity=specific and --reservation=RESERVATION_NAME flags, where RESERVATION_NAME is the name of your TPU reservation. To use Spot VMs, add the --spot flag.

Create a CPU node pool for the Pathways head, which runs the Pathways resource manager and proxy. Name the node pool cpu-np so that the gcluster job submit --pathways command can detect it:
```
gcloud container node-pools create cpu-np \
    --project=${PROJECT} \
    --zone=${ZONE} \
    --cluster=${CLUSTER} \
    --machine-type=${PW_CPU_MACHINE_TYPE} \
    --num-nodes=${NUM_CPU_NODES} \
    --scopes=storage-full,gke-default,cloud-platform \
    --workload-metadata=GCE_METADATA
```
Note: Pathways workloads store temporary files and cached compiled code in a Cloud Storage bucket. Because the node pools in this example use --workload-metadata=GCE_METADATA, Pathways containers access Cloud Storage by using the node service account. By default, this account is the Compute Engine default service account. Grant this service account the Storage Object Admin role (roles/storage.objectAdmin) on the bucket.
Get credentials for the cluster and set the default namespace for kubectl:

Note: If you use a zonal cluster, such as the cluster in this example, specify --zone. If you use a regional cluster, replace --zone=${ZONE} with --region=${REGION}.
```
gcloud container clusters get-credentials ${CLUSTER} \
    --zone=${ZONE} \
    --project=${PROJECT} \
    && kubectl config set-context --current --namespace=default
```

To use the Pathways architecture on your GKE cluster, install the JobSet API:

kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.12.0/manifests.yaml

Verify that the nodes have a STATUS of Ready and that the jobset-controller-manager pod has a STATUS of Running:

kubectl get nodes
kubectl get pods --namespace=jobset-system

The output is similar to the following:

NAME                                      STATUS   ROLES    AGE   VERSION
gke-cluster-cpu-np-12345678-abcd          Ready    <none>   5m    v1.32.2-gke.1475000
gke-cluster-tpu-np-1-12345678-efgh        Ready    <none>   5m    v1.32.2-gke.1475000
...
NAME                                         READY   STATUS    RESTARTS   AGE
jobset-controller-manager-6f7b8c9d4-xyz12    1/1     Running   0          2m

Clean up

To avoid incurring charges to your Google Cloud account for the resources used on this page, delete the cluster when you are finished:

Cluster Toolkit: Destroy the deployment by running the following command:
```
./gcluster destroy DEPLOYMENT_FOLDER_NAME
```
Replace DEPLOYMENT_FOLDER_NAME with the name of your deployment folder.

gcloud CLI: Delete the GKE cluster and its node pools by running the following command:

gcloud container clusters delete ${CLUSTER} \
    --project=${PROJECT} \
    --zone=${ZONE}