Data

Tools

Indexes

Applications

Experiments

Awesome

An open API service indexing awesome lists of open source software.

https://github.com/gke-labs/tpu-operator

A Kubernetes operator designed to automate the lifecycle of TPU slices within self-managed clusters on Google Compute Engine (GCE).
https://github.com/gke-labs/tpu-operator

gce kubernetes tpu tpu-slice

Last synced: about 1 month ago
JSON representation

A Kubernetes operator designed to automate the lifecycle of TPU slices within self-managed clusters on Google Compute Engine (GCE).

Awesome Lists containing this project

README 

          
# tpu-operator



A Kubernetes operator designed to automate the lifecycle of TPU slices within self-managed clusters on Google Compute Engine (GCE).



## Overview: What the Controller Does for You



When running self-managed Kubernetes on Google Compute Engine (GCE), provisioning TPU slices manually requires orchestrating multiple GCE APIs, managing complex metadata for hardware discovery, and bootstrapping nodes.



The TPUNodeGroup controller acts as a Kubernetes-native infrastructure operator that automates the complexities of provisioning TPU slices. Unlike standard GCE VMs, TPU slices involve multiple VMs interconnected with a specialized high-speed network, forming a single, atomic unit for large scale ML workloads.



From your perspective, you simply provide a TPUNodeGroup CR declaring your desired TPU shape. Behind the scenes, the controller will:



1. Abstract GCE Complexity: To properly provision TPUs, the controller automatically creates three critical GCE resources on your behalf:

   * GCE Workload Policy: Responsible for specifying the physical properties of the underlying infrastructure. For multi-host slices, this policy specifies the acceleratorTopology (e.g., 4x4x4), which tells Google's data centers exactly how to physically wire the TPU chips together using high-speed Inter-Chip Interconnects (ICI).

   * Instance Template: Defines the VM configuration (machine type, image, disks, metadata, etc.) used by the Managed Instance Group.

   * Managed Instance Group (MIG): Responsible for actually provisioning and managing the lifecycle of the Virtual Machines. For multi-host slices, users must explicitly configure `targetSizePolicyMode: "BULK"` in the TPUNodeGroup CR. The controller then configures the MIG to use "bulk mode", which guarantees atomic provisioning. This ensures GCE will only create the VMs if it can successfully allocate the entire slice at once, preventing partial, unusable setups.

2. Inject TPU Metadata: Automatically inject TPU configuration metadata needed by the TPU Device Plugin (like kube-labels and accelerator\_topology\_id) into the VMs via the MIG.

3. Secure "Pull-Model" Bootstrapping: To handle long GCE provisioning times (where standard short-lived kubeadm tokens might expire), the controller implements a **Pull Model**. Instead of injecting tokens at VM creation, the controller waits for the VM state to be `RUNNING`, then injects a dynamic, short-lived join token into Instance Metadata. A startup script on the VM polls for this token, ensuring a secure and reliable cluster join even after extended provisioning delays. (Note: This feature is primarily provided to facilitate rapid prototyping and testing. For production environments, it is expected that users will leverage their existing mechanisms for node bootstrapping, such as custom OS images or configuration management tools).

4. Deploy the TPU Device Plugin: The controller automatically deploys a streamlined, open-source TPU device plugin to your cluster as a DaemonSet. This plugin acts as the bridge between the raw TPU hardware and the Kubernetes scheduler. It discovers the attached TPU chips, registers them as schedulable resources (e.g., google.com/tpu), and handles local hardware health monitoring.

5. Handle Graceful Teardown: When the TPUNodeGroup CR is deleted, the controller orchestrates a graceful teardown using a Kubernetes Finalizer. It cordons the nodes, deletes the GCE resources (MIG, Instance Template, and Workload Policy) in reverse order, and removes stale Node objects from the cluster.



*Note: The current release supports v6e and v7x TPU generations.*



---



## Design Philosophy: Reference Bootstrapping



This project is intended as a **Reference Implementation** for experimental and testing purposes. 



**Core Assumption:** We assume that production customers already have mature solutions for their node infrastructure and lifecycle. For those who need a rapid, "turn-key" way to experiment with TPUs on self-managed clusters, we provide a **Reference Bootstrapping** flow using standard Kubernetes mechanics.



## Architecture: The Composite Pattern



This controller is designed using the **Composite Pattern** to maximize modularity and reusability. Instead of a single monolithic state machine, the project is structured as a parent controller that orchestrates three dedicated child Custom Resources:



*   **ResourcePolicy Controller:** Manages GCE Resource Policies (Workload Policies) for ICI networking.

*   **InstanceTemplate Controller:** Manages GCE Instance Templates.

*   **ManagedInstanceGroup Controller:** Manages GCE MIGs, with specific support for bulk allocation.



All infrastructure resources are managed asynchronously via official **Google Cloud Go Client Libraries**. This architecture allows the `TPUNodeGroup` controller to focus solely on high-level workload orchestration and node lifecycle, leaving the low-level GCE "3C" (Compute) resource management to specialized, decoupled sub-controllers.



![TPUNodeGroup Architecture](docs/images/tpu_node_group_architecture.svg)



### The Join-to-Workload Lifecycle (Pull-Model Bootstrapping)



When `bootstrapKubernetes` is enabled, the controller orchestrates a secure, asynchronous **Pull-Model Bootstrapping** flow. This design safely coordinates dynamic token generation and join operations over potentially extended GCE VM booting delays.



The configuration sequence proceeds as follows:



1. **GCE VM Initialization (VM Starts Booting):** The Managed Instance Group (MIG) provisions the VM.

2. **GCE VM (Polling Begins):** The VM's startup script immediately begins polling its own GCE Instance Metadata for the `kubeadm-join-token`.

3. **Controller (RUNNING Detection):** In parallel, the controller detects that the GCE VM state has transitioned to `RUNNING`.

4. **Controller (Token Generation):** The controller generates a unique, short-lived `kubeadm` join token (valid for 1 hour) and creates a corresponding **Bootstrap Secret** in the `kube-system` namespace.

5. **Controller (Metadata Injection):** The controller fetches the cluster's **CA Cert Hash** and injects it along with the **Join Token** and **Control Plane IP** directly into the GCE Instance Metadata for that specific VM.

6. **GCE VM (Poll Succeeds):** The VM's polling loop succeeds and retrieves the injected configuration values.

7. **GCE VM (Pull-Model Join):** The VM executes `kubeadm join` using the retrieved credentials to securely attach itself to the cluster.

8. **Kubernetes (Node Registered):** The K8s API Server registers the node and marks it as `Ready`.

9. **Controller (Node Matching & Labeling):** Once the new worker node appears in the cluster, the controller matches it to the GCE instance using the node's **`ProviderID`** and automatically **patches the Node with TPU labels** (topology, accelerator type, chip count).



The following sequence diagram visualizes this chronological pull-model orchestration:



![Pull-Model Bootstrapping Flow](docs/images/bootstrapping_flow.svg)



---



## Prerequisites



### Configure GCP IAM Permissions



The TPUNodeGroup controller requires permissions to create and manage GCE resources (Instance Templates, Workload Policies, and MIGs) on your behalf.



You can assign the roles/compute.admin role, or for least-privilege environments, ensure the service account has the following specific permissions:



* compute.instanceTemplates.\*

* compute.resourcePolicies.\*

* compute.instanceGroupManagers.\*

* compute.instanceGroups.\*

* compute.instances.\*

* compute.subnetworks.get



### Set Up Kubernetes Control Plane (GCE VM)



If you are setting up a new self-managed cluster, you can initialize a control plane on a GCE VM using `kubeadm`.



1. **Create a GCE VM** (if you don't have one):



```shell

export ZONE=us-central1-c

export SERVICE_ACCOUNT_EMAIL=your-service-account@project.iam.gserviceaccount.com



gcloud compute instances create k8s-control-plane \

    --zone=$ZONE \

    --machine-type=e2-standard-4 \

    --image-family=ubuntu-2404-lts-amd64 \

    --image-project=ubuntu-os-cloud \

    --tags=k8s-control-plane \

    --service-account=$SERVICE_ACCOUNT_EMAIL \

    --scopes=https://www.googleapis.com/auth/cloud-platform

```



2. **Initialize the cluster**: SSH into the VM and run the following commands.

   Ensure `conntrack`, `containerd`, `kubeadm` and a CNI plugin are installed.



```shell

# Initialize the control plane with a specific pod CIDR (Flannel requires 10.244.0.0/16)

sudo kubeadm init --pod-network-cidr=10.244.0.0/16

```



Note: kubeadm is required if enabling `bootstrapKubernetes` in the TPUNodeGroup CR.



### Configure Your Subnet



Ensure your subnet has internet access (e.g., via Cloud NAT) so nodes can download packages during bootstrapping.



### Configure VPC Firewall Rules



Allow API Server access (TCP 6443\) from the TPU nodes and allow internal traffic between the Control Plane and the TPU VMs. Adjust the source-ranges to match your VPC CIDR:



```shell

export VPC_NETWORK=default

export CONTROL_PLANE_CIDR=0.0.0.0/0 # Restrict this in production

export CLUSTER_CIDR=10.128.0.0/9 # Adjust to your VPC CIDR



# Allow API Server access (6443)

gcloud compute firewall-rules create allow-k8s-apiserver \

    --allow tcp:6443 \

    --network=$VPC_NETWORK \

    --source-ranges=$CONTROL_PLANE_CIDR \

    --target-tags=k8s-control-plane



# Allow all internal traffic between nodes (Control Plane <-> TPU VM)

gcloud compute firewall-rules create allow-k8s-internal \

    --allow tcp,udp,icmp,ipip \

    --network=$VPC_NETWORK \

    --source-ranges=$CLUSTER_CIDR \

    --target-tags=k8s-control-plane

```



### Create Your TPU Reservation (Optional)



If you are using a reservation-bound provisioning model, you must reserve the TPU capacity in your target zone.



```shell

export RESERVATION_NAME=my-tpu-reservation

export VM_FAMILY=tpu7x # or tpu6e

export ACCELERATOR_COUNT=4

export HOST_COUNT=2

export ZONE=us-central1-c



gcloud compute reservations create $RESERVATION_NAME \

    --vm-family=$VM_FAMILY \

    --accelerator-count=$ACCELERATOR_COUNT \

    --host-count=$HOST_COUNT \

    --zone=$ZONE

```



##



## Step 1: Build and Deploy the Controller



Before creating the TPUNodeGroup resource, you need to build the controller image, push it to your container registry, and deploy it using the provided manifests.



1. Build and Push the Image:



```shell

export PROJECT_ID=your-project-id



docker build -t gcr.io/$PROJECT_ID/tpunodegroup-controller:latest .

docker push gcr.io/$PROJECT_ID/tpunodegroup-controller:latest

```



2. Update the Deployment Manifest: Open `deploy/kustomization.yaml` and update the `newName` and `newTag` under the `images` section to match the image you just pushed.

3. Install Components: Apply the kustomization to deploy all components (CRDs, Controller, Device Plugin, etc.).



```shell

kubectl apply -k deploy/

```



4. Verify the Controller is Running:



```shell

kubectl get pods -n tpu-node-group

```



*Note: If you are running a single-node cluster (control plane only), you may need to untaint the node to allow the controller to run: `kubectl taint nodes --all node-role.kubernetes.io/control-plane-`*



## Step 2: Create the TPUNodeGroup CR



Once the controller is running, create a TPUNodeGroup resource in your cluster. This tells the controller what capacity and configuration to instantiate.



The controller manages the creation of the Instance Template dynamically based on the `instanceConfig` fields. The `bootstrapKubernetes` block instructs the controller to automatically bootstrap the Kubernetes nodes and join them to the control plane.



Set the required environment variables:



```shell

export PROJECT=my-gcp-project

export ZONE=us-central1-c

export RESERVATION_NAME=my-tpu-reservation

export IMAGE=projects/ubuntu-os-accelerator-images/global/images/ubuntu-accel-2404-amd64-tpu-tpu7x-v20260320

export TOPOLOGY=2x2x2

export NODE_COUNT=2

export CONTROL_PLANE_IP=10.128.0.2

export TPU_NODE_GROUP_NAME=my-tpu-node-group

export NAMESPACE=default

export ACCELERATOR=tpu7x

export MACHINE_TYPE=tpu7x-standard-4t

```



Example CR creation using these variables:



```shell

cat <//`). The controller relies on this `providerID` to correctly identify, match, and label the nodes as they join the cluster. Once matched, the controller will automatically perform node labeling and deploy the TPU device plugin DaemonSet.*



## Step 3: Monitor Provisioning Status



The controller provides real-time readiness feedback in the CRD status. You don't need to check the GCP console; simply inspect the resource in Kubernetes to see its exact lifecycle state.



```shell

kubectl get tpunodegroup $TPU_NODE_GROUP_NAME -o yaml

```



You can view the detailed conditions array to see the health and provisioning phase of your TPU capacity. The conditions dynamically adapt based on your configuration:



* Ready: Set to True only when 100% of the nodes in the slice are healthy, registered, and their ICI links are established.



## Step 4: Schedule Your Workloads



In your workload manifest, you can add Kubernetes node selectors to ensure that your TPU workload is scheduled on the correct TPU machine type and topology. The controller automatically injects these labels into the nodes.



*Note: For JAX workloads, please ensure you use JAX 0.10.0 or newer to avoid the need for manual worker network endpoint metadata configuration.*



Example deployment using node selectors:



```shell

cat <