Using helm charts for kueue installation

This commit refactors the Kueue scheduler installation, migrating it from manually managed static YAML files to the official Kueue Helm chart.

Author: shubpal07

examples/gke-a3-highgpu.yaml

@@ -25,6 +25,7 @@ vars:
   # The following line must be updated for this example to work.
   authorized_cidr: <your-ip-address>/32
   gcp_public_cidrs_access_enabled: false
+  system_node_pool_disk_size_gb: 200
 
 deployment_groups:
 - group: primary
@@ -80,6 +81,9 @@ deployment_groups:
     source: modules/scheduler/gke-cluster
     use: [network1, gpunets, workload_service_account]
     settings:
+      system_node_pool_machine_type: "e2-standard-16"
+      system_node_pool_disk_size_gb: $(vars.system_node_pool_disk_size_gb)
+      system_node_pool_taints: []
       enable_private_endpoint: false  # Allows for access from authorized public IPs
       gcp_public_cidrs_access_enabled: $(vars.gcp_public_cidrs_access_enabled)
       configure_workload_identity_sa: true

examples/gke-a3-megagpu/gke-a3-megagpu.yaml

@@ -54,6 +54,7 @@ vars:
   health_check_schedule: "0 0 * * 0" # Run the health check at 12:00 AM (midnight) every Sunday
 
   permissions_file_staged_path: $(ghpc_stage("./chs-permissions.yaml.tftpl"))
+  system_node_pool_disk_size_gb: 200
   chs_output_bucket_name: chs-result
   chs_pvc_claim_name: chs-output-pvc
   chs_cronjob_rendered_path: $(ghpc_stage("./chs-cronjob.yaml.tftpl"))
@@ -114,6 +115,9 @@ deployment_groups:
     source: modules/scheduler/gke-cluster
     use: [network1, gpunets, workload_service_account]
     settings:
+      system_node_pool_machine_type: "e2-standard-16"
+      system_node_pool_disk_size_gb: $(vars.system_node_pool_disk_size_gb)
+      system_node_pool_taints: []
       enable_private_endpoint: false  # Allows for access from authorized public IPs
       gcp_public_cidrs_access_enabled: $(vars.gcp_public_cidrs_access_enabled)
       configure_workload_identity_sa: true

modules/management/kubectl-apply/README.md

@@ -207,7 +207,7 @@ limitations under the License.
 | <a name="module_install_gib"></a> [install\_gib](#module\_install\_gib) | ./kubectl | n/a |
 | <a name="module_install_gpu_operator"></a> [install\_gpu\_operator](#module\_install\_gpu\_operator) | ./helm_install | n/a |
 | <a name="module_install_jobset"></a> [install\_jobset](#module\_install\_jobset) | ./kubectl | n/a |
-| <a name="module_install_kueue"></a> [install\_kueue](#module\_install\_kueue) | ./kubectl | n/a |
+| <a name="module_install_kueue"></a> [install\_kueue](#module\_install\_kueue) | ./helm_install | n/a |
 | <a name="module_install_nvidia_dra_driver"></a> [install\_nvidia\_dra\_driver](#module\_install\_nvidia\_dra\_driver) | ./helm_install | n/a |
 | <a name="module_kubectl_apply_manifests"></a> [kubectl\_apply\_manifests](#module\_kubectl\_apply\_manifests) | ./kubectl | n/a |
 
@@ -233,7 +233,7 @@ limitations under the License.
 | <a name="input_gke_cluster_exists"></a> [gke\_cluster\_exists](#input\_gke\_cluster\_exists) | A static flag that signals to downstream modules that a cluster has been created. Needed by community/modules/scripts/kubernetes-operations. | `bool` | `false` | no |
 | <a name="input_gpu_operator"></a> [gpu\_operator](#input\_gpu\_operator) | Install [GPU Operator](https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/getting-started.html) which uses the [Kubernetes operator](https://kubernetes.io/docs/concepts/extend-kubernetes/operator/) to automate the management of all NVIDIA software components needed to provision GPU. | <pre>object({<br/>    install = optional(bool, false)<br/>    version = optional(string, "v25.3.0")<br/>  })</pre> | `{}` | no |
 | <a name="input_jobset"></a> [jobset](#input\_jobset) | Install [Jobset](https://github.com/kubernetes-sigs/jobset) which manages a group of K8s [jobs](https://kubernetes.io/docs/concepts/workloads/controllers/job/) as a unit. | <pre>object({<br/>    install = optional(bool, false)<br/>    version = optional(string, "v0.7.2")<br/>  })</pre> | `{}` | no |
-| <a name="input_kueue"></a> [kueue](#input\_kueue) | Install and configure [Kueue](https://kueue.sigs.k8s.io/docs/overview/) workload scheduler. A configuration yaml/template file can be provided with config\_path to be applied right after kueue installation. If a template file provided, its variables can be set to config\_template\_vars. | <pre>object({<br/>    install              = optional(bool, false)<br/>    version              = optional(string, "v0.12.2")<br/>    config_path          = optional(string, null)<br/>    config_template_vars = optional(map(any), null)<br/>  })</pre> | `{}` | no |
+| <a name="input_kueue"></a> [kueue](#input\_kueue) | Install and configure [Kueue](https://kueue.sigs.k8s.io/docs/overview/) workload scheduler. A configuration yaml/template file can be provided with config\_path to be applied right after kueue installation. If a template file provided, its variables can be set to config\_template\_vars. | <pre>object({<br/>    install              = optional(bool, false)<br/>    version              = optional(string, "0.13.3")<br/>    config_path          = optional(string, null)<br/>    config_template_vars = optional(map(any), null)<br/>  })</pre> | `{}` | no |
 | <a name="input_nvidia_dra_driver"></a> [nvidia\_dra\_driver](#input\_nvidia\_dra\_driver) | Installs [Nvidia DRA driver](https://github.com/NVIDIA/k8s-dra-driver-gpu) which supports Dynamic Resource Allocation for NVIDIA GPUs in Kubernetes | <pre>object({<br/>    install = optional(bool, false)<br/>    version = optional(string, "v25.3.0")<br/>  })</pre> | `{}` | no |
 | <a name="input_project_id"></a> [project\_id](#input\_project\_id) | The project ID that hosts the gke cluster. | `string` | n/a | yes |
 

modules/management/kubectl-apply/kueue/kueue-helm-values.yaml

@@ -0,0 +1,30 @@
+# Copyright 2025 "Google LLC"
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# For referencing the original Kueue helm chart values, pull the latest helm chart version
+# `helm pull oci://registry.k8s.io/kueue/charts/kueue --version=0.13.3` (latest helm chart version)
+
+controllerManager:
+  # -- Enables the Topology-Aware Scheduling feature gate.
+  featureGates:
+  - name: TopologyAwareScheduling
+    enabled: true
+
+  # It ensures the Kueue pod can schedule on GKE clusters where the
+  # system node pool uses the default "gke-managed-components" taint.
+  tolerations:
+  - key: "components.gke.io/gke-managed-components"
+    operator: "Equal"
+    value: "true"
+    effect: "NoSchedule"

modules/management/kubectl-apply/main.tf

@@ -53,7 +53,6 @@ locals {
   install_gpu_operator      = try(var.gpu_operator.install, false)
   install_nvidia_dra_driver = try(var.nvidia_dra_driver.install, false)
   install_gib               = try(var.gib.install, false)
-  kueue_install_source      = format("${path.module}/manifests/kueue-%s.yaml", try(var.kueue.version, ""))
   jobset_install_source     = format("${path.module}/manifests/jobset-%s.yaml", try(var.jobset.version, ""))
 }
 
@@ -87,15 +86,21 @@ module "kubectl_apply_manifests" {
 }
 
 module "install_kueue" {
-  source            = "./kubectl"
-  source_path       = local.install_kueue ? local.kueue_install_source : null
-  server_side_apply = true
-  wait_for_rollout  = true
-  depends_on        = [var.gke_cluster_exists]
+  source           = "./helm_install"
+  count            = local.install_kueue ? 1 : 0
+  wait             = false
+  timeout          = 1200
+  release_name     = "kueue"
+  chart_repository = "oci://registry.k8s.io/kueue/charts"
+  chart_name       = "kueue"
+  chart_version    = var.kueue.version
+  namespace        = "kueue-system"
+  create_namespace = true
+  values_yaml = [
+    file("${path.module}/kueue/kueue-helm-values.yaml")
+  ]
 
-  providers = {
-    kubectl = kubectl
-  }
+  depends_on = [var.gke_cluster_exists]
 }
 
 module "configure_kueue" {