feat: promote tpu to ga (#1856)

Author: DrFaust92

README.md

@@ -162,6 +162,7 @@ Then perform the following commands on the root folder:
 | enable\_network\_egress\_export | Whether to enable network egress metering for this cluster. If enabled, a daemonset will be created in the cluster to meter network egress traffic. | `bool` | `false` | no |
 | enable\_resource\_consumption\_export | Whether to enable resource consumption metering on this cluster. When enabled, a table will be created in the resource export BigQuery dataset to store resource consumption data. The resulting table can be joined with the resource usage table or with BigQuery billing export. | `bool` | `true` | no |
 | enable\_shielded\_nodes | Enable Shielded Nodes features on all nodes in this cluster | `bool` | `true` | no |
+| enable\_tpu | Enable Cloud TPU resources in the cluster. WARNING: changing this after cluster creation is destructive! | `bool` | `false` | no |
 | enable\_vertical\_pod\_autoscaling | Vertical Pod Autoscaling automatically adjusts the resources of pods controlled by it | `bool` | `false` | no |
 | filestore\_csi\_driver | The status of the Filestore CSI driver addon, which allows the usage of filestore instance as volumes | `bool` | `false` | no |
 | firewall\_inbound\_ports | List of TCP ports for admission/webhook controllers. Either flag `add_master_webhook_firewall_rules` or `add_cluster_firewall_rules` (also adds egress rules) must be set to `true` for inbound-ports firewall rules to be applied. | `list(string)` | <pre>[<br>  "8443",<br>  "9443",<br>  "15017"<br>]</pre> | no |
@@ -257,6 +258,7 @@ Then perform the following commands on the root folder:
 | region | Cluster region |
 | release\_channel | The release channel of this cluster |
 | service\_account | The service account to default running nodes as if not overridden in `node_pools`. |
+| tpu\_ipv4\_cidr\_block | The IP range in CIDR notation used for the TPUs |
 | type | Cluster type (regional / zonal) |
 | vertical\_pod\_autoscaling\_enabled | Whether vertical pod autoscaling enabled |
 | zones | List of zones in which the cluster resides |

autogen/main/cluster.tf.tmpl

@@ -189,10 +189,9 @@ resource "google_container_cluster" "primary" {
   }
 
   enable_kubernetes_alpha = var.enable_kubernetes_alpha
-
+  enable_tpu                  = var.enable_tpu
   {% if beta_cluster %}
   enable_intranode_visibility = var.enable_intranode_visibility
-  enable_tpu                  = var.enable_tpu
 
   dynamic "pod_security_policy_config" {
     for_each = var.enable_pod_security_policy ? [var.enable_pod_security_policy] : []

autogen/main/firewall.tf.tmpl

@@ -57,7 +57,6 @@ resource "google_compute_firewall" "intra_egress" {
 }
 
 
-{% if beta_cluster %}
 /******************************************
   Allow egress to the TPU IPv4 CIDR block
 
@@ -95,8 +94,6 @@ resource "google_compute_firewall" "tpu_egress" {
   {% endif %}
 }
 
-
-{% endif %}
 /******************************************
   Allow GKE master to hit non 443 ports for
   Webhooks/Admission Controllers

autogen/main/outputs.tf.tmpl

@@ -171,6 +171,11 @@ output "identity_namespace" {
   ]
 }
 
+output "tpu_ipv4_cidr_block" {
+  description = "The IP range in CIDR notation used for the TPUs"
+  value       = var.enable_tpu ? google_container_cluster.primary.tpu_ipv4_cidr_block : null
+}
+
 {% if autopilot_cluster != true %}
 output "mesh_certificates_config" {
   description = "Mesh certificates configuration"
@@ -228,9 +233,4 @@ output "identity_service_enabled" {
   description = "Whether Identity Service is enabled"
   value       = local.cluster_pod_security_policy_enabled
 }
-
-output "tpu_ipv4_cidr_block" {
-  description = "The IP range in CIDR notation used for the TPUs"
-  value       = var.enable_tpu ? google_container_cluster.primary.tpu_ipv4_cidr_block : null
-}
 {% endif %}

autogen/main/variables.tf.tmpl

@@ -600,13 +600,12 @@ variable "deletion_protection" {
   default     = true
 }
 
-{% if beta_cluster %}
 variable "enable_tpu" {
   type        = bool
   description = "Enable Cloud TPU resources in the cluster. WARNING: changing this after cluster creation is destructive!"
   default     = false
 }
-{% endif %}
+
 {% if autopilot_cluster != true %}
 variable "network_policy" {
   type        = bool

cluster.tf

@@ -146,7 +146,7 @@ resource "google_container_cluster" "primary" {
   }
 
   enable_kubernetes_alpha = var.enable_kubernetes_alpha
-
+  enable_tpu              = var.enable_tpu
   dynamic "master_authorized_networks_config" {
     for_each = local.master_authorized_networks_config
     content {

firewall.tf

@@ -55,6 +55,41 @@ resource "google_compute_firewall" "intra_egress" {
 }
 
 
+/******************************************
+  Allow egress to the TPU IPv4 CIDR block
+
+  This rule is defined separately from the
+  intra_egress rule above since it requires
+  an output from the google_container_cluster
+  resource.
+
+  https://github.com/terraform-google-modules/terraform-google-kubernetes-engine/issues/1124
+ *****************************************/
+resource "google_compute_firewall" "tpu_egress" {
+  count       = var.add_cluster_firewall_rules && var.enable_tpu ? 1 : 0
+  name        = "gke-${substr(var.name, 0, min(36, length(var.name)))}-tpu-egress"
+  description = "Managed by terraform gke module: Allow pods to communicate with TPUs"
+  project     = local.network_project_id
+  network     = var.network
+  priority    = var.firewall_priority
+  direction   = "EGRESS"
+
+  target_tags        = [local.cluster_network_tag]
+  destination_ranges = [google_container_cluster.primary.tpu_ipv4_cidr_block]
+
+  # Allow all possible protocols
+  allow { protocol = "tcp" }
+  allow { protocol = "udp" }
+  allow { protocol = "icmp" }
+  allow { protocol = "sctp" }
+  allow { protocol = "esp" }
+  allow { protocol = "ah" }
+
+  depends_on = [
+    google_container_cluster.primary,
+  ]
+}
+
 /******************************************
   Allow GKE master to hit non 443 ports for
   Webhooks/Admission Controllers

modules/beta-autopilot-private-cluster/firewall.tf

@@ -84,7 +84,6 @@ resource "google_compute_firewall" "tpu_egress" {
 
 }
 
-
 /******************************************
   Allow GKE master to hit non 443 ports for
   Webhooks/Admission Controllers

modules/beta-autopilot-private-cluster/outputs.tf

@@ -142,6 +142,11 @@ output "identity_namespace" {
   ]
 }
 
+output "tpu_ipv4_cidr_block" {
+  description = "The IP range in CIDR notation used for the TPUs"
+  value       = var.enable_tpu ? google_container_cluster.primary.tpu_ipv4_cidr_block : null
+}
+
 
 
 output "master_ipv4_cidr_block" {
@@ -183,8 +188,3 @@ output "identity_service_enabled" {
   description = "Whether Identity Service is enabled"
   value       = local.cluster_pod_security_policy_enabled
 }
-
-output "tpu_ipv4_cidr_block" {
-  description = "The IP range in CIDR notation used for the TPUs"
-  value       = var.enable_tpu ? google_container_cluster.primary.tpu_ipv4_cidr_block : null
-}

modules/beta-autopilot-private-cluster/variables.tf

@@ -433,6 +433,7 @@ variable "enable_tpu" {
   description = "Enable Cloud TPU resources in the cluster. WARNING: changing this after cluster creation is destructive!"
   default     = false
 }
+
 variable "database_encryption" {
   description = "Application-layer Secrets Encryption settings. The object format is {state = string, key_name = string}. Valid values of state are: \"ENCRYPTED\"; \"DECRYPTED\". key_name is the name of a CloudKMS key."
   type        = list(object({ state = string, key_name = string }))