TPU Provisioner: Add support for v6e & cross project reservations

tpu-provisioner/go.mod

@@ -1,6 +1,6 @@
 module github.com/GoogleCloudPlatform/ai-on-gke/tpu-provisioner
 
-go 1.22.0
+go 1.23.0
 
 require (
 	cloud.google.com/go/compute/metadata v0.3.0

tpu-provisioner/internal/cloud/common.go

@@ -26,6 +26,11 @@ const (
 
 	// AnnotationCopyLabels is a comma-separated list of labels to copy from the Pod to the node pool config (Nodes).
 	AnnotationCopyLabels = "tpu-provisioner.cloud.google.com/copy-labels"
+	// AnnotationAdditionalNodeNetworks is a comma-separated list of additional networks and subnets to attach to the node pool.
+	// Format: "<network-name>:<subnet-name>, ..."
+	AnnotationAdditionalNodeNetworks = "tpu-provisioner.cloud.google.com/additional-node-networks"
+	// AnnotatationServiceAccount is the GCP service account to use for the node pool.
+	AnnotationNodeServiceAccount = "tpu-provisioner.cloud.google.com/node-service-account"
 
 	EventNodePoolCreationStarted   = "NodePoolCreationStarted"
 	EventNodePoolCreationSucceeded = "NodePoolCreationSucceeded"

tpu-provisioner/internal/cloud/gke.go

@@ -42,6 +42,7 @@ const (
 	V4PodSliceAccelerator  = "tpu-v4-podslice"
 	V5ePodSliceAccelerator = "tpu-v5-lite-podslice"
 	V5pPodSliceAccelerator = "tpu-v5p-slice"
+	V6eSliceAccelerator    = "tpu-v6e-slice"
 
 	// Resource type labels
 	GoogleTPUResource = "google.com/tpu"
@@ -355,10 +356,40 @@ func (g *GKE) nodePoolForPod(name string, p *corev1.Pod) (*containerv1beta1.Node
 		}
 	}
 
+	var networkConfig *containerv1beta1.NodeNetworkConfig
+	var additionalNodeNetworks []*containerv1beta1.AdditionalNodeNetworkConfig
+	// additional-node-networks: "vpc1:subnet1, vpc2:subnet2"
+	for _, pair := range strings.Split(getAnnotation(p, AnnotationAdditionalNodeNetworks), ",") {
+		pair = strings.TrimSpace(pair)
+		if pair == "" {
+			continue
+		}
+
+		netAndSubnet := strings.SplitN(pair, ":", 2)
+		if len(netAndSubnet) != 2 {
+			return nil, fmt.Errorf("invalid additional network annotation: %v", pair)
+		}
+
+		additionalNodeNetworks = append(additionalNodeNetworks, &containerv1beta1.AdditionalNodeNetworkConfig{
+			Network:    strings.TrimSpace(netAndSubnet[0]),
+			Subnetwork: strings.TrimSpace(netAndSubnet[1]),
+		})
+	}
+	if len(additionalNodeNetworks) > 0 {
+		networkConfig = &containerv1beta1.NodeNetworkConfig{
+			AdditionalNodeNetworkConfigs: additionalNodeNetworks,
+		}
+	}
+
+	nodeServiceAccount := g.ClusterContext.NodeServiceAccount
+	if sa, ok := p.Annotations[AnnotationNodeServiceAccount]; ok {
+		nodeServiceAccount = sa
+	}
+
 	return &containerv1beta1.NodePool{
 		Name: name,
 		Config: &containerv1beta1.NodeConfig{
-			ServiceAccount: g.ClusterContext.NodeServiceAccount,
+			ServiceAccount: nodeServiceAccount,
 			ShieldedInstanceConfig: &containerv1beta1.ShieldedInstanceConfig{
 				EnableIntegrityMonitoring: true,
 				EnableSecureBoot:          g.ClusterContext.NodeSecureBoot,
@@ -387,6 +418,7 @@ func (g *GKE) nodePoolForPod(name string, p *corev1.Pod) (*containerv1beta1.Node
 			MaxSurge: 1,
 		},
 		MaxPodsConstraint: &containerv1beta1.MaxPodsConstraint{MaxPodsPerNode: maxPodsPerNode},
+		NetworkConfig:     networkConfig,
 	}, nil
 }
 
@@ -438,7 +470,7 @@ func tpuTopologyToNodeCount(accelerator, topo string) (int, error) {
 	switch accelerator {
 	case V4PodSliceAccelerator, V5pPodSliceAccelerator:
 		expectedDims = 3
-	case V5ePodSliceAccelerator:
+	case V5ePodSliceAccelerator, V6eSliceAccelerator:
 		expectedDims = 2
 	default:
 		return 0, fmt.Errorf("invalid accelerator: %v", accelerator)
@@ -475,6 +507,8 @@ func tpuMachineType(accel string, tpuRequest int) (string, error) {
 		return fmt.Sprintf("ct5lp-hightpu-%vt", tpuRequest), nil
 	case V5pPodSliceAccelerator: // v5p
 		return fmt.Sprintf("ct5p-hightpu-%vt", tpuRequest), nil
+	case V6eSliceAccelerator: // v6e
+		return fmt.Sprintf("ct6e-standard-%vt", tpuRequest), nil
 	}
 
 	return "", fmt.Errorf("invalid accelerator: %v", accel)

tpu-provisioner/internal/cloud/gke_test.go

@@ -67,6 +67,16 @@ func Test_tpuTopologyToNodeCount(t *testing.T) {
 			topo:  "not-a-topo",
 			err:   true,
 		},
+		{
+			accel: "tpu-v6e-slice",
+			topo:  "16x16",
+			count: 64,
+		},
+		{
+			accel: "tpu-v6e-slice",
+			topo:  "1x1x1",
+			err:   true,
+		},
 	}
 
 	for _, c := range cases {