Update cluster as regional and update ip_cidr_range, enable vpc-nativ…

…e traffic routing and the example for v5e.

PiperOrigin-RevId: 566429751

tools/kubernetes/terraform/examples/v4/main.tf

@@ -12,4 +12,4 @@ module "tpu-gke" {
   region               = var.region
   tpu_node_pools       = var.tpu_node_pools
   maintenance_interval = var.maintenance_interval
-}
+}

tools/kubernetes/terraform/examples/v4/outputs.tf

@@ -21,4 +21,4 @@ output "kubernetes_cluster_host" {
 output "nodepool_tpu_topology" {
   value       = module.tpu-gke.nodepool_tpu_topology
   description = "GKE TPU topology"
-}
+}

tools/kubernetes/terraform/examples/v4/terraform.tfvars

@@ -6,4 +6,4 @@ tpu_node_pools = [{
   node_count   = 2
   machine_type = "ct4p-hightpu-4t"
   topology     = "2x2x2"
-  }]
+}]

tools/kubernetes/terraform/examples/v5e/main.tf

@@ -12,4 +12,4 @@ module "tpu-gke" {
   region               = var.region
   tpu_node_pools       = var.tpu_node_pools
   maintenance_interval = var.maintenance_interval
-}
+}

tools/kubernetes/terraform/examples/v5e/outputs.tf

@@ -18,7 +18,7 @@ output "kubernetes_cluster_host" {
   description = "GKE Cluster Host"
 }
 
-output "nodepool_tpu_topology" {
-  value       = module.tpu-gke.nodepool_tpu_topology
-  description = "GKE TPU topology"
+output "placement_policy_names" {
+  value       = module.tpu-gke.placement_policy_names
+  description = "GKE TPU Placement Policy Names"
 }

tools/kubernetes/terraform/examples/v5e/terraform.tfvars

@@ -1,15 +1,101 @@
-project_id           = "project-id"
-resource_name_prefix = "tpu-v5lite-test"
+project_id           = "project_id"
+resource_name_prefix = "tpu-v5e-test"
 region               = "us-east5"
 tpu_node_pools = [{
   zone         = "us-east5-b"
-  node_count   = 1
+  node_count   = 64
   machine_type = "ct5lp-hightpu-4t"
-  topology     = "2x2"
+  topology     = "16x16"
+  policy       = "sb-compact-4a"
   }, {
   zone         = "us-east5-b"
-  node_count   = 1
+  node_count   = 64
   machine_type = "ct5lp-hightpu-4t"
-  topology     = "2x2"
+  topology     = "16x16"
+  policy       = "sb-compact-4a"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4a"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4a"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4b"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4b"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4b"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4b"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4c"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4c"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4c"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4c"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4d"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4d"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4d"
+  }, {
+  zone         = "us-east5-b"
+  node_count   = 64
+  machine_type = "ct5lp-hightpu-4t"
+  topology     = "16x16"
+  policy       = "sb-compact-4d"
 }]
-maintenance_interval = "PERIODIC"
+maintenance_interval = "PERIODIC"

tools/kubernetes/terraform/module/main.tf

@@ -36,7 +36,7 @@ resource "google_compute_subnetwork" "subnet" {
   name          = "${var.resource_name_prefix}-subnet"
   region        = var.region
   network       = google_compute_network.vpc.name
-  ip_cidr_range = "10.10.0.0/24"
+  ip_cidr_range = "10.10.0.0/19"
 }
 
 resource "google_container_cluster" "tpu_cluster" {
@@ -48,6 +48,12 @@ resource "google_container_cluster" "tpu_cluster" {
   # node pool and immediately delete it.
   remove_default_node_pool = true
   initial_node_count       = 1
+  networking_mode          = "VPC_NATIVE"
+  ip_allocation_policy {
+    cluster_ipv4_cidr_block  = "/14"
+    services_ipv4_cidr_block = "/20"
+  }
+  default_max_pods_per_node = 50
 
   release_channel {
     channel = "UNSPECIFIED"
@@ -84,6 +90,7 @@ resource "google_container_node_pool" "multihost_tpu" {
     oauth_scopes = [
       "https://www.googleapis.com/auth/logging.write",
       "https://www.googleapis.com/auth/monitoring",
+      "https://www.googleapis.com/auth/cloud-platform",
     ]
     host_maintenance_policy {
       maintenance_interval = var.maintenance_interval
@@ -98,15 +105,15 @@ resource "google_container_node_pool" "multihost_tpu" {
       enabled = true
     }
 
-    image_type = "COS_CONTAINERD"
+    image_type   = "COS_CONTAINERD"
     machine_type = var.tpu_node_pools[count.index].machine_type
     tags         = ["gke-node"]
     metadata = {
       disable-legacy-endpoints = "true"
     }
   }
   placement_policy {
-    type         = "COMPACT"
-    tpu_topology = var.tpu_node_pools[count.index].topology
+    type        = "COMPACT"
+    policy_name = var.tpu_node_pools[count.index].policy
   }
 }

tools/kubernetes/terraform/module/outputs.tf

@@ -18,7 +18,9 @@ output "kubernetes_cluster_host" {
   description = "GKE Cluster Host"
 }
 
-output "nodepool_tpu_topology" {
-  value       = flatten(google_container_node_pool.multihost_tpu[*].placement_policy[0].tpu_topology)
-  description = "GKE TPU topology"
-}
+output "placement_policy_names" {
+  value = flatten([
+    google_container_node_pool.multihost_tpu[*].placement_policy[0].policy_name
+  ])
+  description = "GKE TPU Placement Policy Names"
+}

tools/kubernetes/terraform/module/terraform.tfvars

@@ -1,6 +1,7 @@
 project_id           = "project-id"
 resource_name_prefix = "tpu-test"
 region               = "us-central2"
+location             = "us-central2-b"
 tpu_node_pools = [{
   zone         = "us-central2-b"
   node_count   = 4
@@ -17,4 +18,4 @@ tpu_node_pools = [{
   machine_type = "ct4p-hightpu-4t"
   topology     = "2x2x2"
 }]
-maintenance_interval = "AS_NEEDED"
+maintenance_interval = "AS_NEEDED"

tools/kubernetes/terraform/module/variables.tf

@@ -34,10 +34,11 @@ variable "tpu_node_pools" {
     node_count   = number,
     machine_type = string,
     topology     = string,
+    policy       = string,
   }))
 }
 
 variable "maintenance_interval" {
-  default = "AS_NEEDED"
+  default     = "AS_NEEDED"
   description = "maintenance interval for TPU machines."
-}
+}