Adding changes to tpu-definition and gke-node-pool for workload_policy

Author: agrawalkhushi18

modules/compute/gke-node-pool/main.tf

@@ -113,7 +113,7 @@ resource "google_container_node_pool" "node_pool" {
     content {
       type         = var.placement_policy.type
       policy_name  = var.placement_policy.name
-      tpu_topology = var.placement_policy.tpu_topology
+      tpu_topology = module.tpu.is_tpu ? var.placement_policy.tpu_topology : null
     }
   }
 

modules/compute/resource-policy/README.md

@@ -78,5 +78,5 @@ No modules.
 
 | Name | Description |
 |------|-------------|
-| <a name="output_placement_policy"></a> [placement\_policy](#output\_placement\_policy) | Group placement policy to use for placing VMs or GKE nodes placement. `COMPACT` is the only supported value for `type` currently. `name` is the name of the placement policy.<br/>It is assumed that the specified policy exists. To create a placement policy refer to https://cloud.google.com/sdk/gcloud/reference/compute/resource-policies/create/group-placement.<br/>Note: Placement policies have the [following](https://cloud.google.com/compute/docs/instances/placement-policies-overview#restrictions-compact-policies) restrictions. |
+| <a name="output_placement_policy"></a> [placement\_policy](#output\_placement\_policy) | Group placement policy to use for placing VMs or GKE nodes placement. `COMPACT` is the only supported value for `type` currently. `name` is the name of the placement policy.<br/>It is assumed that the specified policy exists. To create a placement policy refer to https://cloud.google.com/sdk/gcloud/reference/compute/resource-policies/create/group-placement.<br/>Note: Placement policies have the [following](https://cloud.google.com/compute/docs/instances/placement-policies-overview#restrictions-compact-policies) restrictions.<br/>The value `tpu_topology` is only used for TPU node pools. The `gke-node-pool` module ensures it is configured appropriately for only TPUs during placement policy mapping. |
 <!-- END OF PRE-COMMIT-TERRAFORM DOCS HOOK -->

modules/compute/resource-policy/outputs.tf

@@ -19,10 +19,12 @@ output "placement_policy" {
   Group placement policy to use for placing VMs or GKE nodes placement. `COMPACT` is the only supported value for `type` currently. `name` is the name of the placement policy.
   It is assumed that the specified policy exists. To create a placement policy refer to https://cloud.google.com/sdk/gcloud/reference/compute/resource-policies/create/group-placement.
   Note: Placement policies have the [following](https://cloud.google.com/compute/docs/instances/placement-policies-overview#restrictions-compact-policies) restrictions.
+  The value `tpu_topology` is only used for TPU node pools. The `gke-node-pool` module ensures it is configured appropriately for only TPUs during placement policy mapping.
   EOT
 
   value = {
-    type = (var.group_placement_max_distance > 0 || var.workload_policy.type != null) ? "COMPACT" : null
-    name = (var.group_placement_max_distance > 0 || var.workload_policy.type != null) ? local.name : null
+    type         = (var.group_placement_max_distance > 0 || var.workload_policy.type != null) ? "COMPACT" : null
+    name         = (var.group_placement_max_distance > 0 || var.workload_policy.type != null) ? local.name : null
+    tpu_topology = (var.workload_policy.type != null) ? var.workload_policy.accelerator_topology : null
   }
 }

modules/internal/tpu-definition/main.tf

@@ -15,16 +15,16 @@
 */
 
 locals {
-  # Safely check for the existence of a TPU topology.
-  is_tpu = try(var.placement_policy.tpu_topology, null) != null && try(var.placement_policy.tpu_topology, "") != ""
+  # Determine if this is a TPU node pool by checking if the machine_type exists in our authoritative map of TPU machine types.
+  is_tpu = contains(keys(local.tpu_chip_count_map), var.machine_type)
 
   tpu_taint = local.is_tpu ? [{
     key    = "google.com/tpu"
     value  = "present"
     effect = "NO_SCHEDULE"
   }] : []
 
-  # Exhaustive map of machine prefixes to GKE accelerator labels.
+  # Map of machine prefixes to GKE accelerator labels.
   tpu_accelerator_map = {
     "ct4p"  = "tpu-v4-podslice"      # TPU v4
     "ct5lp" = "tpu-v5-lite-podslice" # TPU v5e
@@ -33,7 +33,7 @@ locals {
   }
 
   # Map specific GCE machine types to the number of TPU chips per node (VM).
-  # Based on public Google Cloud TPU documentation.
+  # The machine-type map must be updated to reflect new TPU releases with reference to public documentation: https://docs.cloud.google.com/tpu/docs/intro-to-tpu
   tpu_chip_count_map = {
     # v4 - ct4p
     "ct4p-hightpu-4t" = 4