Initial implementation of tpu_config in container cluster (#6130) (#4390)

Co-authored-by: Shuya Ma <[email protected]>
Signed-off-by: Modular Magician <[email protected]>

Co-authored-by: Shuya Ma <[email protected]>

Author: modular-magician

.changelog/6130.txt

@@ -0,0 +1,3 @@
+```release-note:enhancement
+container: added `tpu_config` to `google_container_cluster` (beta only)
+```

google-beta/resource_container_cluster.go

@@ -521,11 +521,43 @@ func resourceContainerCluster() *schema.Resource {
 			},
 
 			"enable_tpu": {
-				Type:        schema.TypeBool,
+				Type:          schema.TypeBool,
+				Optional:      true,
+				ForceNew:      true,
+				Description:   `Whether to enable Cloud TPU resources in this cluster.`,
+				ConflictsWith: []string{"tpu_config"},
+				Computed:      true,
+				// TODO: deprecate when tpu_config is correctly returned by the API
+				// Deprecated: "Deprecated in favor of tpu_config",
+			},
+
+			"tpu_config": {
+				Type:        schema.TypeList,
 				Optional:    true,
-				ForceNew:    true,
-				Description: `Whether to enable Cloud TPU resources in this cluster.`,
-				Default:     false,
+				Computed:    true,
+				MaxItems:    1,
+				Description: `TPU configuration for the cluster.`,
+				Elem: &schema.Resource{
+					Schema: map[string]*schema.Schema{
+						"enabled": {
+							Type:        schema.TypeBool,
+							Required:    true,
+							ForceNew:    true,
+							Description: `Whether Cloud TPU integration is enabled or not`,
+						},
+						"ipv4_cidr_block": {
+							Type:        schema.TypeString,
+							Computed:    true,
+							Description: `IPv4 CIDR block reserved for Cloud TPU in the VPC.`,
+						},
+						"use_service_networking": {
+							Type:        schema.TypeBool,
+							Optional:    true,
+							ForceNew:    true,
+							Description: `Whether to use service networking for Cloud TPU or not`,
+						},
+					},
+				},
 			},
 
 			"enable_legacy_abac": {
@@ -1606,6 +1638,10 @@ func resourceContainerClusterCreate(d *schema.ResourceData, meta interface{}) er
 		cluster.IdentityServiceConfig = expandIdentityServiceConfig(v)
 	}
 
+	if v, ok := d.GetOk("tpu_config"); ok {
+		cluster.TpuConfig = expandContainerClusterTpuConfig(v)
+	}
+
 	if v, ok := d.GetOk("resource_usage_export_config"); ok {
 		cluster.ResourceUsageExportConfig = expandResourceUsageExportConfig(v)
 	}
@@ -3633,6 +3669,19 @@ func expandMonitoringConfig(configured interface{}) *container.MonitoringConfig
 	return mc
 }
 
+func expandContainerClusterTpuConfig(configured interface{}) *container.TpuConfig {
+	l := configured.([]interface{})
+	if len(l) == 0 || l[0] == nil {
+		return nil
+	}
+
+	config := l[0].(map[string]interface{})
+	return &container.TpuConfig{
+		Enabled:              config["enabled"].(bool),
+		UseServiceNetworking: config["use_service_networking"].(bool),
+	}
+}
+
 func flattenNotificationConfig(c *container.NotificationConfig) []map[string]interface{} {
 	if c == nil {
 		return nil

google-beta/resource_container_cluster_test.go

@@ -2527,6 +2527,30 @@ func TestAccContainerCluster_withDNSConfig(t *testing.T) {
 	})
 }
 
+func TestAccContainerCluster_withTPUConfig(t *testing.T) {
+	t.Parallel()
+
+	clusterName := fmt.Sprintf("tf-test-cluster-%s", randString(t, 10))
+	containerNetName := fmt.Sprintf("tf-test-container-net-%s", randString(t, 10))
+	vcrTest(t, resource.TestCase{
+		PreCheck:     func() { testAccPreCheck(t) },
+		Providers:    testAccProviders,
+		CheckDestroy: testAccCheckContainerClusterDestroyProducer(t),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccContainerCluster_withTPUConfig(containerNetName, clusterName),
+			},
+			{
+				ResourceName:      "google_container_cluster.with_tpu_config",
+				ImportState:       true,
+				ImportStateVerify: true,
+				// TODO: remove when tpu_config can be read from the API
+				ImportStateVerifyIgnore: []string{"tpu_config"},
+			},
+		},
+	})
+}
+
 func testAccContainerCluster_masterAuthorizedNetworksDisabled(t *testing.T, resource_name string) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
 		rs, ok := s.RootModule().Resources[resource_name]
@@ -5271,3 +5295,59 @@ resource "google_container_cluster" "primary" {
 }
 `, name, name, name)
 }
+
+func testAccContainerCluster_withTPUConfig(network, cluster string) string {
+	return fmt.Sprintf(`
+resource "google_compute_network" "container_network" {
+	name                    = "%s"
+	auto_create_subnetworks = false
+}
+	  
+resource "google_compute_subnetwork" "container_subnetwork" {
+	name                     = google_compute_network.container_network.name
+	network                  = google_compute_network.container_network.name
+	ip_cidr_range            = "10.0.36.0/24"
+	region                   = "us-central1"
+	private_ip_google_access = true
+
+	secondary_ip_range {
+		range_name    = "pod"
+		ip_cidr_range = "10.0.0.0/19"
+	}
+
+	secondary_ip_range {
+		range_name    = "svc"
+		ip_cidr_range = "10.0.32.0/22"
+	}
+}
+	
+resource "google_container_cluster" "with_tpu_config" {
+	name               = "%s"
+	location           = "us-central1-a"
+	initial_node_count = 1
+	
+	
+	tpu_config {
+		enabled                = true
+		use_service_networking = true
+	}
+	
+	network         = google_compute_network.container_network.name
+	subnetwork      = google_compute_subnetwork.container_subnetwork.name
+	networking_mode = "VPC_NATIVE"
+	
+	private_cluster_config {
+		enable_private_endpoint = true
+		enable_private_nodes    = true
+		master_ipv4_cidr_block  = "10.42.0.0/28"
+	}
+	master_authorized_networks_config {
+	}
+	
+	ip_allocation_policy {
+		cluster_secondary_range_name  = google_compute_subnetwork.container_subnetwork.secondary_ip_range[0].range_name
+		services_secondary_range_name = google_compute_subnetwork.container_subnetwork.secondary_ip_range[1].range_name
+	}
+}
+`, network, cluster)
+}