Skip to content

Feat: Automatically derive TPU node counts based on topology and machine type#5386

Merged
SwarnaBharathiMantena merged 7 commits into
GoogleCloudPlatform:developfrom
SwarnaBharathiMantena:swarnabm/align_static_node_count_tpu
Mar 30, 2026
Merged

Feat: Automatically derive TPU node counts based on topology and machine type#5386
SwarnaBharathiMantena merged 7 commits into
GoogleCloudPlatform:developfrom
SwarnaBharathiMantena:swarnabm/align_static_node_count_tpu

Conversation

@SwarnaBharathiMantena
Copy link
Copy Markdown
Contributor

@SwarnaBharathiMantena SwarnaBharathiMantena commented Mar 23, 2026 •
edited
Loading

Description

This PR improves the usability and developer experience for provisioning TPU workloads in the Cluster Toolkit by eliminating the requirement for users to manually compute and specify static_node_count.

Key Features & Changes

  1. Automatic Node Derivation: Natively computes the exact amount of required computing nodes (static_node_count) directly from the provided machine_type and tpu_topology.
  2. Hardware Standardization: Implements generic mapping logic for chipsPerVM across various TPU generations (v4, v5p, v5e, v6e, v7x) as well as gracefully parsing legacy string literals (like v5litepod-16).
  3. Placement Policy Injection: Automatically configures and injects a COMPACT placement policy for derived multi-node TPU cluster setups.
  4. Backward Compatibility: Preserves standard behavior for non-TPU resources. CPU and GPU parsing remains unaffected with settings passing through as usual.
  5. Testing & Coverage: Introduces robust, table-driven unit test coverage in pkg/config/hardware_test.go to validate topology extraction, the core evaluation logic in expandHardwareSettings, exact chip calculation math, and placement policy placement.

Context

Abstracting this logic directly into gcluster resolves the previous gap between high-level accelerator shorthand and declarative infrastructure, ensuring configurations are fail-fast and user YAML files are significantly simplified.

@SwarnaBharathiMantena SwarnaBharathiMantena added the release-improvements Added to release notes under the "Improvements" heading. label Mar 23, 2026
@gemini-code-assist
Copy link
Copy Markdown
Contributor

gemini-code-assist Bot commented Mar 23, 2026

Summary of Changes

Hello, 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 introduces an automated mechanism to determine the static_node_count for TPU blueprints. By inferring this value from the provided machine_type and tpu_topology, it streamlines the configuration process, reduces manual errors, and ensures consistent hardware provisioning for TPU-based modules, while maintaining compatibility with other hardware types.

Highlights

  • Automated TPU Node Calculation: The static_node_count variable for TPU blueprints is now automatically computed based on machine_type and tpu_topology inputs, simplifying configuration.
  • New Hardware Logic Module: A new file, pkg/config/hardware.go, was introduced to encapsulate the logic for extracting TPU topology, evaluating machine types, and calculating the number of TPU nodes.
  • Placement Policy Injection: For multi-node TPU configurations (where static_node_count is greater than 1), a 'COMPACT' placement policy is automatically injected, including the tpu_topology string.
  • Backward Compatibility: The changes ensure that existing GPU and CPU blueprints remain unaffected and continue to function as before, only applying to TPUs when static_node_count is not explicitly set.
  • Comprehensive Unit Tests: New unit tests were added in pkg/config/hardware_test.go to validate the calculateTPUNodes function across various TPU configurations and error scenarios, ensuring robustness.
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 <command> or @gemini-code-assist <command>. 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 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 counter productive. 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[bot]
gemini-code-assist Bot reviewed Mar 23, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces new functionality to automatically infer and expand hardware settings for TPU modules, including calculating static_node_count and injecting a compact placement policy based on machine type and topology. This logic is integrated into the module expansion process. While the calculateTPUNodes function is well-tested, there is a lack of unit tests for expandHardwareSettings, extractTopology, and injectCompactPlacementPolicy. Additionally, there are opportunities to improve code clarity and efficiency by replacing a magic number with a named constant and moving the familyDefaults map to a package-level variable.

Comment thread pkg/config/hardware_test.go Outdated
Comment thread pkg/config/hardware.go Outdated
Comment thread pkg/config/hardware.go Outdated
@SwarnaBharathiMantena SwarnaBharathiMantena marked this pull request as ready for review March 24, 2026 06:27
@SwarnaBharathiMantena SwarnaBharathiMantena changed the title Abstract TPU chipsPerVM logic Feat: Automatically derive TPU node counts based on topology and machine type Mar 24, 2026
Neelabh94
Neelabh94 reviewed Mar 26, 2026
View reviewed changes
Comment thread pkg/config/hardware.go Outdated
@SwarnaBharathiMantena SwarnaBharathiMantena merged commit 4cea3b8 into GoogleCloudPlatform:develop Mar 30, 2026
15 of 72 checks passed
@agrawalkhushi18 agrawalkhushi18 mentioned this pull request Apr 28, 2026
Merged
@agrawalkhushi18 agrawalkhushi18 mentioned this pull request May 5, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kvenkatachala333 kvenkatachala333 kvenkatachala333 approved these changes

@samskillman samskillman Awaiting requested review from samskillman

@cboneti cboneti Awaiting requested review from cboneti

@Neelabh94 Neelabh94 Awaiting requested review from Neelabh94

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

release-improvements Added to release notes under the "Improvements" heading.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@SwarnaBharathiMantena @Neelabh94 @kvenkatachala333