From f0766dd9acc4c1c783240c2eb658f1eae48158f4 Mon Sep 17 00:00:00 2001
From: annuay <annuay@google.com>
Date: Mon, 10 Feb 2025 08:48:03 +0000
Subject: [PATCH] Fix guest accelerators for GKE

---
 modules/internal/gpu-definition/README.md |  2 +-
 modules/internal/gpu-definition/main.tf   | 10 +++++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/modules/internal/gpu-definition/README.md b/modules/internal/gpu-definition/README.md
index 29a87cab78..16ea533010 100644
--- a/modules/internal/gpu-definition/README.md
+++ b/modules/internal/gpu-definition/README.md
@@ -35,7 +35,7 @@ No resources.
 
 | Name | Description | Type | Default | Required |
 |------|-------------|------|---------|:--------:|
-| <a name="input_guest_accelerator"></a> [guest\_accelerator](#input\_guest\_accelerator) | List of the type and count of accelerator cards attached to the instance. | <pre>list(object({<br/>    type  = string,<br/>    count = number<br/>  }))</pre> | `[]` | no |
+| <a name="input_guest_accelerator"></a> [guest\_accelerator](#input\_guest\_accelerator) | List of the type and count of accelerator cards attached to the instance. | <pre>list(object({<br/>    type  = string<br/>    count = number<br/>    gpu_driver_installation_config = optional(object({<br/>      gpu_driver_version = string<br/>    }), { gpu_driver_version = "DEFAULT" })<br/>    gpu_partition_size = optional(string)<br/>    gpu_sharing_config = optional(object({<br/>      gpu_sharing_strategy       = string<br/>      max_shared_clients_per_gpu = number<br/>    }))<br/>  }))</pre> | `[]` | no |
 | <a name="input_machine_type"></a> [machine\_type](#input\_machine\_type) | Machine type to use for the instance creation | `string` | n/a | yes |
 
 ## Outputs
diff --git a/modules/internal/gpu-definition/main.tf b/modules/internal/gpu-definition/main.tf
index bc66442e5e..9181ae1aa9 100644
--- a/modules/internal/gpu-definition/main.tf
+++ b/modules/internal/gpu-definition/main.tf
@@ -22,8 +22,16 @@ variable "machine_type" {
 variable "guest_accelerator" {
   description = "List of the type and count of accelerator cards attached to the instance."
   type = list(object({
-    type  = string,
+    type  = string
     count = number
+    gpu_driver_installation_config = optional(object({
+      gpu_driver_version = string
+    }), { gpu_driver_version = "DEFAULT" })
+    gpu_partition_size = optional(string)
+    gpu_sharing_config = optional(object({
+      gpu_sharing_strategy       = string
+      max_shared_clients_per_gpu = number
+    }))
   }))
   default  = []
   nullable = false