feat: Dynamic Slicing: Install Kueue Slice Controller and set default resources - #5794
Conversation
…nership collision
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request integrates the Kueue slice-controller into the cluster toolkit to support GKE Dynamic Slicing for TPUs. By leveraging a raw YAML manifest deployment strategy, the changes ensure the controller is installed without Helm ownership issues, maintaining clean namespace management. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request introduces the installation of a slice controller under the kubectl-apply management module, including its Kubernetes manifests and a Terraform module definition. The review feedback highlights several critical improvements: pinning the container image to a stable version from an official registry instead of using a personal Docker Hub repository with the :latest tag, adding a Terraform lifecycle precondition to ensure Kueue is installed when slicing is enabled, lowering the excessively high CPU and memory resource requests/limits to prevent scheduling failures, and explicitly setting the Helm release namespace to kube-system for consistency.
- Added a precondition in variables.tf to explicitly fail if kueue.enable_dynamic_slicing_for_tpus is true but kueue.install is false. - Wrapped manifest.source expressions in try() within main.tf to fix terraform validation crashes when manifest.source is null.
- Converted slice-controller.yaml to a Terraform template. - Added variables for cpu/memory limits and requests inside the kueue variable block, keeping the high default values from upstream docs.
- When dynamic slicing is enabled, automatically default Kueue and Jobset controller resources to the higher values recommended in the official GKE dynamic slicing documentation. - Kueue defaults: 3 replicas, 16 cpu, 64Gi memory. - Jobset defaults: 4 cpu, 16Gi memory. - User-provided values will still override these defaults.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds support for dynamic TPU slicing by introducing a slice controller module, configuring resource allocations for Kueue and Jobset, and adding validation preconditions. The review feedback suggests addressing a potential startup race condition by making the slice controller depend on the Jobset installation, adding a validation precondition to ensure Jobset is installed when dynamic slicing is enabled, and pinning the slice controller container image to a specific version instead of using the ":latest" tag.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces the Kueue slice controller to support dynamic slicing for TPUs, adding a new deployment template, updating the main configuration to conditionally install the controller, and adding input variables and preconditions. The review feedback highlights critical security improvements, including restricting cluster-wide secret access in the RBAC configuration to a namespaced Role, and parameterizing the hardcoded personal Docker Hub image with the ":latest" tag to allow configuring a pinned, official version. Additionally, the feedback recommends keeping the new variables nullable and using try() blocks to safely access attributes and prevent plan-time crashes.
|
Regarding the high priority comments in the PR, they are related to the manifest that is publicly available in the docs page. I don't think we should be fixing that? |
|
@SwarnaBharathiMantena @Neelabh94 please run the e2e tests |
|
/gcbrun |
e07d51c
into
GoogleCloudPlatform:develop
Description
This PR integrates the Kueue Slice Controller into the
kubectl-applymodule to support GKE dynamic slicing for TPUs.The controller is automatically deployed when
kueue.enable_dynamic_slicing_for_tpusis enabled.Key Changes
slice-controller.yamlmanifest directly from the GKE Dynamic Slicing Documentation. It is natively applied via the internalraw-config-chart.slice-controller-systemnamespace is managed explicitly within the manifest rather than through Helm release tracking, avoiding tracking collisions.slice-controller.yaml.tftpl). Users can now override the massive default resource requirements via thekueuevariable block if deploying on smaller node pools. It defaults to the upstream recommended limits (12000mcpu,32Gimemory).jobsetandkueuecontroller resources to the higher capacities recommended in the GKE docs (Jobset: 4 cpu / 16Gi memory, Kueue: 16 cpu / 64Gi memory / 3 replicas) unless explicitly overridden by the user.lifecycleprecondition that explicitly throws an error if a user attempts to enable dynamic slicing while Kueue installation is disabled.