TPU Provisioner: Add confidential disk config (#1005)

Author: nstogner

tpu-provisioner/cmd/main.go

@@ -85,6 +85,10 @@ func main() {
 		GCPNodeSecureBoot         bool     `envconfig:"GCP_NODE_SECURE_BOOT" default:"true"`
 		GCPNodeAdditionalNetworks string   `envconfig:"GCP_NODE_ADDITIONAL_NETWORKS" default:""`
 
+		GCPNodeDiskType            string `envconfig:"GCP_NODE_DISK_TYPE"`
+		GCPNodeConfidentialStorage bool   `envconfig:"GCP_NODE_CONFIDENTIAL_STORAGE"`
+		GCPNodeBootDiskKMSKey      string `envconfig:"GCP_NODE_BOOT_DISK_KMS_KEY"`
+
 		// GCPForceOnDemand forces the controller to create nodes on demand, even if
 		// the Pod requests a reservation or spot.
 		GCPForceOnDemand bool `envconfig:"GCP_FORCE_ON_DEMAND" default:"false"`
@@ -202,17 +206,20 @@ func main() {
 		provider = &cloud.GKE{
 			Service: containers,
 			ClusterContext: cloud.GKEContext{
-				ProjectID:              cfg.GCPProjectID,
-				ClusterLocation:        cfg.GCPClusterLocation,
-				Cluster:                cfg.GCPCluster,
-				NodeZone:               cfg.GCPZone,
-				NodeServiceAccount:     cfg.GCPNodeServiceAccount,
-				NodeAdditionalNetworks: cfg.GCPNodeAdditionalNetworks,
-				NodeSecondaryDisk:      cfg.GCPNodeSecondaryDisk,
-				NodeTags:               cfg.GCPNodeTags,
-				PodToNodeLabels:        cfg.GCPPodToNodeLabels,
-				NodeSecureBoot:         cfg.GCPNodeSecureBoot,
-				ForceOnDemand:          cfg.GCPForceOnDemand,
+				ProjectID:               cfg.GCPProjectID,
+				ClusterLocation:         cfg.GCPClusterLocation,
+				Cluster:                 cfg.GCPCluster,
+				NodeZone:                cfg.GCPZone,
+				NodeServiceAccount:      cfg.GCPNodeServiceAccount,
+				NodeAdditionalNetworks:  cfg.GCPNodeAdditionalNetworks,
+				NodeSecondaryDisk:       cfg.GCPNodeSecondaryDisk,
+				NodeTags:                cfg.GCPNodeTags,
+				NodeDiskType:            cfg.GCPNodeDiskType,
+				NodeConfidentialStorage: cfg.GCPNodeConfidentialStorage,
+				NodeBootDiskKMSKey:      cfg.GCPNodeBootDiskKMSKey,
+				PodToNodeLabels:         cfg.GCPPodToNodeLabels,
+				NodeSecureBoot:          cfg.GCPNodeSecureBoot,
+				ForceOnDemand:           cfg.GCPForceOnDemand,
 			},
 			Recorder: mgr.GetEventRecorderFor("tpu-provisioner"),
 		}

tpu-provisioner/internal/cloud/gke.go

@@ -405,6 +405,11 @@ func (g *GKE) nodePoolForPod(name string, p *corev1.Pod) (*containerv1beta1.Node
 		placementPolicy.Type = "COMPACT"
 	}
 
+	var diskType string
+	if g.ClusterContext.NodeDiskType != "" {
+		diskType = g.ClusterContext.NodeDiskType
+	}
+
 	return &containerv1beta1.NodePool{
 		Name: name,
 		Config: &containerv1beta1.NodeConfig{
@@ -416,12 +421,15 @@ func (g *GKE) nodePoolForPod(name string, p *corev1.Pod) (*containerv1beta1.Node
 			Tags: g.ClusterContext.NodeTags,
 			// NOTE: vendor/ was manually updated to include the field because
 			// it was not currently available at the time of writing:
-			SecondaryBootDisks:  secondaryDisks,
-			MachineType:         machineType,
-			ReservationAffinity: reservation,
-			Labels:              labels,
-			Spot:                spot,
-			Taints:              taints,
+			SecondaryBootDisks:        secondaryDisks,
+			MachineType:               machineType,
+			ReservationAffinity:       reservation,
+			Labels:                    labels,
+			Spot:                      spot,
+			Taints:                    taints,
+			BootDiskKmsKey:            g.ClusterContext.NodeBootDiskKMSKey,
+			DiskType:                  diskType,
+			EnableConfidentialStorage: g.ClusterContext.NodeConfidentialStorage,
 		},
 		InitialNodeCount: int64(nodeCount),
 		Locations:        []string{g.ClusterContext.NodeZone},

tpu-provisioner/internal/cloud/gke_context.go

@@ -3,14 +3,17 @@ package cloud
 import "fmt"
 
 type GKEContext struct {
-	ProjectID              string
-	ClusterLocation        string
-	Cluster                string
-	NodeZone               string
-	NodeServiceAccount     string
-	NodeAdditionalNetworks string
-	NodeSecondaryDisk      string
-	NodeTags               []string
+	ProjectID               string
+	ClusterLocation         string
+	Cluster                 string
+	NodeZone                string
+	NodeServiceAccount      string
+	NodeAdditionalNetworks  string
+	NodeSecondaryDisk       string
+	NodeTags                []string
+	NodeDiskType            string
+	NodeConfidentialStorage bool
+	NodeBootDiskKMSKey      string
 	// PodToNodeLabels is a list of key=value pairs that will be copied from the Pod
 	// to the Node.
 	PodToNodeLabels []string

tpu-provisioner/internal/cloud/gke_test.go

@@ -691,6 +691,38 @@ func TestNodePoolForPod(t *testing.T) {
 				},
 			},
 		},
+		{
+			desc: "confidential disk configured in cluster context",
+			gkeContext: GKEContext{
+				NodeConfidentialStorage: true,
+				NodeDiskType:            "hyperdisk-balanced",
+				NodeBootDiskKMSKey:      "my-kms-key",
+			},
+			want: &containerv1beta1.NodePool{
+				Config: &container.NodeConfig{
+					Labels: map[string]string{
+						"google.com/nodepool-manager":                 "tpu-provisioner",
+						"google.com/tpu-provisioner-jobset-name":      "jobset-test",
+						"google.com/tpu-provisioner-jobset-namespace": "default",
+						"google.com/tpu-provisioner-parent-kind":      "job",
+						"google.com/tpu-provisioner-parent-name":      "jobset-test-job-1-0",
+						"google.com/tpu-provisioner-parent-namespace": "default",
+					},
+					MachineType:               "ct5p-hightpu-4t",
+					ShieldedInstanceConfig:    &container.ShieldedInstanceConfig{EnableIntegrityMonitoring: true},
+					EnableConfidentialStorage: true,
+					BootDiskKmsKey:            "my-kms-key",
+					DiskType:                  "hyperdisk-balanced",
+				},
+				InitialNodeCount:  512,
+				Locations:         []string{""},
+				Management:        &container.NodeManagement{AutoRepair: true, AutoUpgrade: false},
+				MaxPodsConstraint: &container.MaxPodsConstraint{MaxPodsPerNode: 15},
+				Name:              "test-pool",
+				PlacementPolicy:   &container.PlacementPolicy{TpuTopology: "8x16x16", Type: "COMPACT"},
+				UpgradeSettings:   &container.UpgradeSettings{MaxSurge: 1},
+			},
+		},
 	}
 	for _, tc := range tests {
 		t.Run(tc.desc, func(t *testing.T) {