Merge pull request GoogleCloudPlatform#4765 from shubpal07/shubham/jobset-on-helm

Migration of jobset from static manifests to helm chart and upgrading version to 0.10.1

Author: shubpal07

community/examples/xpk-n2-filestore/xpk-n2-filestore.yaml

@@ -32,7 +32,7 @@ vars:
   slice_cpu_count: 64
   xpk_version: v0.8.0
   kueue_version: v0.11.1
-  jobset_version: v0.7.2
+  jobset_version: 0.10.1
   kjob_version: f775609365df47bdc3e7c10290d8efae7f512464
 
 deployment_groups:

modules/management/kubectl-apply/README.md

@@ -206,7 +206,7 @@ limitations under the License.
 | <a name="module_configure_kueue"></a> [configure\_kueue](#module\_configure\_kueue) | ./kubectl | n/a |
 | <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_jobset"></a> [install\_jobset](#module\_install\_jobset) | ./helm_install | 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 |
@@ -232,7 +232,7 @@ limitations under the License.
 | <a name="input_gib"></a> [gib](#input\_gib) | Install the NCCL gIB plugin | <pre>object({<br/>    install = bool<br/>    path    = string<br/>    template_vars = object({<br/>      image   = optional(string, "us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib")<br/>      version = string<br/>      node_affinity = optional(any, {<br/>        requiredDuringSchedulingIgnoredDuringExecution = {<br/>          nodeSelectorTerms = [{<br/>            matchExpressions = [{<br/>              key      = "cloud.google.com/gke-gpu",<br/>              operator = "In",<br/>              values   = ["true"]<br/>            }]<br/>          }]<br/>        }<br/>      })<br/>      accelerator_count = number<br/>    })<br/>  })</pre> | <pre>{<br/>  "install": false,<br/>  "path": "",<br/>  "template_vars": {<br/>    "accelerator_count": 0,<br/>    "version": ""<br/>  }<br/>}</pre> | no |
 | <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_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, "0.10.1")<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/jobset/jobset-helm-values.yaml

@@ -0,0 +1,25 @@
+# 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 jobset helm chart values, pull the latest jobset chart version
+# `helm pull oci://registry.k8s.io/jobset/charts/jobset --version=0.10.1` (latest helm chart version)
+
+controller:
+  # It ensures the Jobset pod(s) can be scheduled 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)
-  jobset_install_source     = format("${path.module}/manifests/jobset-%s.yaml", try(var.jobset.version, ""))
 }
 
 data "http" "manifest_from_url" {
@@ -118,15 +117,20 @@ module "configure_kueue" {
 }
 
 module "install_jobset" {
-  source            = "./kubectl"
-  source_path       = local.install_jobset ? local.jobset_install_source : null
-  server_side_apply = true
-  wait_for_rollout  = true
-  depends_on        = [var.gke_cluster_exists, module.configure_kueue]
-
-  providers = {
-    kubectl = kubectl
-  }
+  source           = "./helm_install"
+  count            = local.install_jobset ? 1 : 0
+  wait             = false
+  timeout          = 1200
+  release_name     = "jobset"
+  chart_repository = "oci://registry.k8s.io/jobset/charts"
+  chart_name       = "jobset"
+  chart_version    = var.jobset.version
+  namespace        = "jobset-system"
+  create_namespace = true
+  values_yaml = [
+    file("${path.module}/jobset/jobset-helm-values.yaml")
+  ]
+  depends_on = [var.gke_cluster_exists, module.configure_kueue]
 }
 
 module "install_nvidia_dra_driver" {
@@ -178,9 +182,9 @@ module "install_nvidia_dra_driver" {
             operator: Equal
             value: present
             effect: NoSchedule
-          - key: kubernetes.io/arch 
-            operator: Equal 
-            value: arm64 
+          - key: kubernetes.io/arch
+            operator: Equal
+            value: arm64
             effect: NoSchedule
 
       EOF
@@ -223,17 +227,17 @@ module "install_gpu_operator" {
             operator: Equal
             value: present
             effect: NoSchedule
-          - key: kubernetes.io/arch 
-            operator: Equal 
-            value: arm64 
+          - key: kubernetes.io/arch
+            operator: Equal
+            value: arm64
             effect: NoSchedule
 
       node-feature-discovery:
         worker:
           tolerations:
-            - key: kubernetes.io/arch 
-              operator: Equal 
-              value: arm64 
+            - key: kubernetes.io/arch
+              operator: Equal
+              value: arm64
               effect: NoSchedule
             - key: "node-role.kubernetes.io/master"
               operator: "Equal"