From 9432d5b8f73dab4e9ae8200c991a89df5517cfbc Mon Sep 17 00:00:00 2001
From: umeshkumhar <umeshkumhar@google.com>
Date: Thu, 12 Oct 2023 00:35:23 +0530
Subject: [PATCH] add support for gpu_driver_installation_config

---
 autogen/main/cluster.tf.tmpl                       |  7 +++++++
 cluster.tf                                         | 14 ++++++++++++++
 .../beta-private-cluster-update-variant/cluster.tf | 14 ++++++++++++++
 modules/beta-private-cluster/cluster.tf            | 14 ++++++++++++++
 .../beta-public-cluster-update-variant/cluster.tf  | 14 ++++++++++++++
 modules/beta-public-cluster/cluster.tf             | 14 ++++++++++++++
 modules/private-cluster-update-variant/cluster.tf  | 14 ++++++++++++++
 modules/private-cluster/cluster.tf                 | 14 ++++++++++++++
 8 files changed, 105 insertions(+)

diff --git a/autogen/main/cluster.tf.tmpl b/autogen/main/cluster.tf.tmpl
index 714b90fb8e..7e4b02ccd9 100644
--- a/autogen/main/cluster.tf.tmpl
+++ b/autogen/main/cluster.tf.tmpl
@@ -846,6 +846,13 @@ resource "google_container_node_pool" "windows_pools" {
         type  = lookup(each.value, "accelerator_type", "")
         count = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
diff --git a/cluster.tf b/cluster.tf
index 7c8b05cdaa..b2d4524254 100644
--- a/cluster.tf
+++ b/cluster.tf
@@ -528,6 +528,13 @@ resource "google_container_node_pool" "pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+        
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
@@ -716,6 +723,13 @@ resource "google_container_node_pool" "windows_pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
diff --git a/modules/beta-private-cluster-update-variant/cluster.tf b/modules/beta-private-cluster-update-variant/cluster.tf
index 29770ab8f6..30563cfef5 100644
--- a/modules/beta-private-cluster-update-variant/cluster.tf
+++ b/modules/beta-private-cluster-update-variant/cluster.tf
@@ -734,6 +734,13 @@ resource "google_container_node_pool" "pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
@@ -962,6 +969,13 @@ resource "google_container_node_pool" "windows_pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
diff --git a/modules/beta-private-cluster/cluster.tf b/modules/beta-private-cluster/cluster.tf
index 1ddb1a76e8..86854f2224 100644
--- a/modules/beta-private-cluster/cluster.tf
+++ b/modules/beta-private-cluster/cluster.tf
@@ -640,6 +640,13 @@ resource "google_container_node_pool" "pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
@@ -867,6 +874,13 @@ resource "google_container_node_pool" "windows_pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
diff --git a/modules/beta-public-cluster-update-variant/cluster.tf b/modules/beta-public-cluster-update-variant/cluster.tf
index d35a15d495..41d0b578f4 100644
--- a/modules/beta-public-cluster-update-variant/cluster.tf
+++ b/modules/beta-public-cluster-update-variant/cluster.tf
@@ -715,6 +715,13 @@ resource "google_container_node_pool" "pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
@@ -943,6 +950,13 @@ resource "google_container_node_pool" "windows_pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
diff --git a/modules/beta-public-cluster/cluster.tf b/modules/beta-public-cluster/cluster.tf
index 4857f57520..fb75975479 100644
--- a/modules/beta-public-cluster/cluster.tf
+++ b/modules/beta-public-cluster/cluster.tf
@@ -621,6 +621,13 @@ resource "google_container_node_pool" "pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
@@ -848,6 +855,13 @@ resource "google_container_node_pool" "windows_pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
diff --git a/modules/private-cluster-update-variant/cluster.tf b/modules/private-cluster-update-variant/cluster.tf
index 191c99ec50..cdbfb7eb78 100644
--- a/modules/private-cluster-update-variant/cluster.tf
+++ b/modules/private-cluster-update-variant/cluster.tf
@@ -641,6 +641,13 @@ resource "google_container_node_pool" "pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
@@ -830,6 +837,13 @@ resource "google_container_node_pool" "windows_pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
diff --git a/modules/private-cluster/cluster.tf b/modules/private-cluster/cluster.tf
index c2fa9906dc..bd996adaeb 100644
--- a/modules/private-cluster/cluster.tf
+++ b/modules/private-cluster/cluster.tf
@@ -547,6 +547,13 @@ resource "google_container_node_pool" "pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }
 
@@ -735,6 +742,13 @@ resource "google_container_node_pool" "windows_pools" {
         type               = lookup(each.value, "accelerator_type", "")
         count              = lookup(each.value, "accelerator_count", 0)
         gpu_partition_size = lookup(each.value, "gpu_partition_size", null)
+
+        dynamic "gpu_driver_installation_config" {
+          for_each = lookup(each.value, "gpu_driver_version", "") != "" ? [1] : []
+          content {
+            gpu_driver_version   = lookup(each.value, "gpu_driver_version", "")
+          }
+        }
       }
     }