fix(job): remove topology labels from Kueue ResourceFlavor nodeLabels

[Neelabh94]

Description

Fixes an issue where gcluster injected node-pool specific physical topology labels (cloud.google.com/gke-tpu-topology: ) into Kueue ResourceFlavor.spec.nodeLabels.

Because Kueue enforces ResourceFlavor.spec.nodeLabels across all workloads using that flavor, Kueue constrained multi-cube and multi-slice TPU workloads (e.g. 4x4x8 and 4x8x8) to a single 16-node block (4x4x4), capping pod allocation at 16 nodes.

Fix

1. Removed cloud.google.com/gke-tpu-topology label injection in resolveAccelerators (pkg/orchestrator/gke/gke_job_orchestrator.go).
2. Added defensive label filtering in renderResourceFlavor (pkg/orchestrator/gke/infra_manager.go) to prevent topology keys from entering Kueue ResourceFlavor specs.

Verification

• Unit tests (go test ./...) passed cleanly.
• Verified on live cluster (bodaborg-super-tpu7x-y6k):
  • 4x4x8 workload launched and ran on 32 nodes (previously capped at 16).
  • 4x8x8 workload launched and ran on 64 nodes (previously capped at 16).

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 addresses a critical limitation in how Kueue ResourceFlavors handle TPU-specific node labels. By removing the automatic injection of topology-related labels and implementing a robust filtering layer, the changes ensure that Kueue does not inadvertently constrain large-scale TPU workloads to single node blocks, thereby enabling proper allocation for multi-cube and multi-slice configurations.

Highlights

• Removed Topology Label Injection: Stopped injecting 'cloud.google.com/gke-tpu-topology' labels into Kueue ResourceFlavor nodeLabels within the GKE job orchestrator.
• Added Defensive Label Filtering: Implemented a filtering mechanism in the infrastructure manager to prevent TPU-specific topology, slice, and partition labels from being included in ResourceFlavor specifications.
• Improved Workload Scalability: Resolved an issue where overly restrictive labels were capping multi-cube and multi-slice TPU workloads, allowing them to scale beyond 16 nodes.

Using Gemini Code Assist

The 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 /gemini or @gemini-code-assist . Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request removes the automatic addition of the TPU topology label to GKE node labels and filters out TPU-related labels (topology, slice, and partition) from the rendered resource flavor. The review feedback suggests replacing the hardcoded "cloud.google.com/gke-tpu-topology" string literal with the existing tpuTopologyLabel constant to improve maintainability.

[Neelabh94]

/gemini review

[gemini-code-assist]

Code Review

This pull request removes the TPU topology label assignment from the GKE orchestrator and introduces filtering logic in renderResourceFlavor to exclude TPU topology, slice, and partition labels from the Kueue ResourceFlavor specification. The review feedback correctly identifies a lack of unit test coverage for this new filtering logic and suggests adding a comprehensive unit test to verify that blocked labels are filtered out while allowed labels are preserved.

[Neelabh94]

/gemini review

[gemini-code-assist]

Code Review

This pull request removes the TPU topology label assignment from the GKE job orchestrator's accelerator resolution and introduces a filtering mechanism in renderResourceFlavor to exclude TPU topology, slice, and partition labels from the rendered resource flavor. A new unit test has been added to verify this filtering behavior. I have no feedback to provide as the changes are clean and well-tested.

[Neelabh94]

/gemini review

[gemini-code-assist]

Code Review

This pull request updates the GKE job orchestrator to filter out TPU topology, slice, and partition labels when rendering resource flavors, and adjusts the exclusive topology configuration to account for static slicing in addition to dynamic slicing. Relevant unit tests have been updated and added to verify this filtering behavior. I have no additional feedback to provide.