feat(gke): Unify A3 Mega blueprint with dynamic consumption models and align daily integration tests

[ishitachail]

Overview

This PR introduces a unified GKE A3 Mega blueprint that natively supports all A3 Mega consumption models (On-Demand, Spot, DWS Flex Start, DWS Queued Provisioning and Specific Reservations) via configurable blueprint variables, and aligns the daily integration tests to work cleanly with the new architecture.

---

Why these changes were made

• Consolidated Consumption Models: Previously, deploying GKE A3 Mega across different consumption models required modifying core blueprint structures or maintaining separate blueprints. This change unifies all consumption options into a single, highly configurable blueprint.
• Dynamic Autoscaler Support for DWS: When using DWS Flex Start with Queued Provisioning, Kueue must use a dedicated queues configuration (dws-queues.yaml.tftpl) so pods trigger GKE's ProvisioningRequest API and dynamically provision nodes, rather than admitting pods against static GPU quota.

---

What changed

1. Unified A3 Mega Blueprint (gke-a3-megagpu.yaml & gke-a3-megagpu-deployment.yaml)

• Configurable Consumption Variables: Added spot, reservation_affinity, enable_flex_start, enable_queued_provisioning, auto_repair, and autoscaling_total_min_nodes to vars: with safe defaults (NO_RESERVATION, spot: false, 0 min nodes).
• Dynamic Conditional Overrides: Configured inline HCL expressions in a3_megagpu_pool to automatically override static_node_count: null and auto_repair: false whenever Flex Start is active, satisfying module preconditions without manual user toggling:
  • static_node_count: "$((vars.enable_flex_start) ? null : vars.static_node_count)"
  • auto_repair: "$((vars.enable_flex_start) ? false : vars.auto_repair)"
• Deployment Templates: Provided commented-out sections in deployment.yaml for each consumption model to serve as an in-line user template. Additionally, under Option 2 (DWS Flex Start + Queued Provisioning), we include the commented-out kueue_configuration_path variable defaulting to $(ghpc_stage("../dws-queues/dws-queues.yaml.tftpl")) along with an inline comment explaining that users can modify this variable if they want to use a custom Kueue configuration file instead of the default DWS queues path.

---

2. Documentation & Workload Samples (examples/dws-sample-workloads/)

• README.md: Added consumption options setup notes, a new Run a Job section, updated 3-mode Verify NCCL Performance instructions (Standard, DWS Flex Start, DWS+Queue), and cleaned up resource links.
• Sample Manifests (examples/dws-sample-workloads/): Added Sample job files dws and dws+queue. Files under this folder are machine independent.

---

3. Daily Integration Test Alignment (tools/cloud-build/daily-tests/)

• Spot Test (gke-a3-megagpu-onspot.yaml & .yml):
  • Removed legacy sed commands that rewrote reservation_affinity: and deleted variables.
  • Added spot: "{{ enable_spot | default(true) }}" to cli_deployment_vars: so Spot provisioning is enabled natively via the blueprint variable.
  • Added a clean sed command to inject placement_policy: { type: "COMPACT" } into a3_megagpu_pool, ensuring GKE attaches GCE topology block/subblock labels required by Kueue Topology Aware Scheduling (TAS).
• Reservation Test (gke-a3-megagpu.yaml & .yml):
  • Added a clean sed command to configure reservation_affinity: for specific reservation testing (a3mega-reservation-0).
  • Removed obsolete reservation: CLI variables from cli_deployment_vars: to avoid test_deployment_variable_not_used validator failures.

---

Verification Performed

• Blueprint Expansion (gcluster create):
  • Verified local blueprint expansion with --vars spot=true, confirming correct Terraform generation (spot = var.spot and placement_policy: { type: COMPACT } in primary/main.tf).
  • Verified non-spot reservation affinity generation (reservation_affinity = var.reservation_affinity).
• Validator Compatibility: Confirmed that removing undeclared CLI variables resolves blueprint validator failures.
• YAML Syntax: Validated that all modified YAML/YML files parse cleanly without syntax errors.

Submission Checklist

NOTE: Community submissions can take up to 2 weeks to be reviewed.

Please take the following actions before submitting this pull request.

• Fork your PR branch from the Toolkit "develop" branch (not main)
• Test all changes with pre-commit in a local branch #
• Confirm that "make tests" passes all tests
• Add or modify unit tests to cover code changes
• Ensure that unit test coverage remains above 80%
• Update all applicable documentation
• Follow Cluster Toolkit Contribution guidelines #

[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 refactors the GKE A3 Mega blueprint to provide a unified, flexible architecture that supports various consumption models natively. By introducing configurable variables and dynamic HCL logic, the change eliminates the need for maintaining separate blueprints or using fragile shell-based test overrides. These improvements simplify deployment management and ensure better compatibility with GKE's advanced provisioning features.

Highlights

• Unified GKE A3 Mega Blueprint: Consolidated multiple consumption models (On-Demand, Spot, DWS Flex Start, Queued Provisioning, and Specific Reservations) into a single, highly configurable blueprint using new variables.
• Dynamic Configuration Logic: Implemented HCL expressions to dynamically toggle node pool settings like auto-repair and static node counts based on the selected consumption model.
• Integration Test Cleanup: Removed brittle 'sed' hacks from daily integration tests, replacing them with native blueprint variable support for cleaner and more reliable test execution.
• DWS Queued Provisioning Support: Added support for GKE's ProvisioningRequest API via a new dedicated Kueue configuration template, triggered automatically when queued provisioning is enabled.

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 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 introduces support for multiple consumption models (such as DWS Flex Start, Queued Provisioning, Spot, and Specific Reservations) in the GKE A3 MegaGPU example blueprint. It adds a new Kueue configuration template (dws-queues.yaml.tftpl) and updates the deployment configuration and daily tests to handle these models dynamically. Feedback on the changes highlights a fragile sed command in the daily tests that can result in malformed YAML, and suggests using single quotes instead of escaped double quotes in the blueprint to improve readability.

[gemini-code-assist]

The sed command matching /^ *spot: / is extremely fragile because it matches both the top-level spot: false variable under vars: and the spot: $(vars.spot) setting under a3_megagpu_pool. This results in appending placement_policy under vars:, which creates malformed and invalid YAML. Instead, target a unique line within the a3_megagpu_pool settings, such as machine_type: a3-megagpu-8g.

    sed -i -e '/machine_type: a3-megagpu-8g/a\\      placement_policy:\\n        type: COMPACT' $${EXAMPLE_BP}

References

1. Prefer exact string matching in sed commands within build and test scripts to maintain consistency with existing codebase patterns, rather than introducing complex regular expressions.

[gemini-code-assist]

Using escaped double quotes \" inside a double-quoted YAML string is hard to read and maintain. You can use single quotes ' inside the double-quoted string for the blueprint expression literals to improve readability and avoid backslash escaping.

  kueue_configuration_path: "$((vars.enable_queued_provisioning) ? ghpc_stage('./dws-queues.yaml.tftpl') : ghpc_stage('./kueue-configuration.yaml.tftpl'))"

[aslam-quad]

/gcbrun

[ishitachail]

/gemini review

[gemini-code-assist]

Code Review

This pull request adds support for GKE A3 Mega GPU consumption options, specifically Dynamic Workload Scheduler (DWS) Flex Start and DWS Flex Start with Queued Provisioning, including sample workloads, Kueue configurations, and updated documentation. Feedback on the changes suggests setting the JobSet configurations to a suspended state (suspend: true) to ensure proper integration with Kueue, and extracting complex inline shell scripts from the initContainers into separate script files to avoid runtime package installations and nested shells.

[agrawalkhushi18]

/gcbrun

[agrawalkhushi18]

Probably, put a double ## here so that when the user un-comments this block, the relevant settings are enabled and not the line information.

[agrawalkhushi18]

I believe we can remove this block now since we are providing a comprehensive description of consumptions above.