[Neelabh94]

Summary

Replaces the raw substring check (requirePathwaysMTCVolume) in GeneratePathwaysManifest with structural Kubernetes PodSpec volume and mount validation (validatePathwaysMTCManifest and validatePodSpecVolumes).

Problems Addressed

1. String-based validation of the rendered Pathways manifest misses broken volume wiring: Checking only for the substring multitier-checkpoint.csi.storage.gke.io in the rendered pathways_jobset.tmpl manifest passes even if:
   • a YAML comment contains the driver string while the CSI volume is omitted,
   • the cache CSI volume is declared in spec.volumes without a matching volumeMounts entry at --gke-mtc-ramdisk-dir on the worker container (or colocated-python-sidecar when --pathways-colocated-python-sidecar-image is set), or
   • a stale custom template still declares the MTC CSI volume on pathways-head, which runs on CPU nodes and gets stuck in FailedMount.
2. ValidateJobSetManifest only checked restartPolicy on JobSet and Job manifests: Custom templates with duplicate volume names, duplicate container volumeMounts paths, or volumeMounts referencing undefined volumes passed client-side manifest validation.

Changes

• pkg/orchestrator/gke/manifest_generator.go:
  • Unmarshal corev1.PodSpec in ValidateJobSetManifest and run validatePodSpecVolumes across all JobSet replicatedJobs and Job documents to catch duplicate volume names, undefined volume mounts, and duplicate mountPaths per container.
  • Replace requirePathwaysMTCVolume in GeneratePathwaysManifest with validatePathwaysMTCManifest(jobSet, opts):
    • Reject pathways-head if it declares a multitier-checkpoint.csi.storage.gke.io CSI volume.
    • Require a worker replicatedJob whose PodSpec declares a CSI volume with driver: multitier-checkpoint.csi.storage.gke.io and mounts that volume at opts.GKEMTCRamdiskDirectory in pathways-worker (and in colocated-python-sidecar when --pathways-colocated-python-sidecar-image is set).
• pkg/orchestrator/gke/manifest_generator_test.go:
  • Expand TestGeneratePathwaysManifest_MTCRejectsStaleCustomTemplate to cover missing worker replicatedJob, comment-only driver mention, MTC volume declared on pathways-head, missing volumeMounts on pathways-worker, mismatched mountPath, and missing volumeMounts on colocated-python-sidecar.
  • Add TestValidateJobSetManifest_PodSpecVolumes covering duplicate volumes and undefined volume references in containers and initContainers, plus TestValidateJobSetManifest_DuplicateMountPathNormalized and TestValidatePathwaysMTCManifest_CleansMountPath for normalized mountPath comparison.

Testing

• go test ./pkg/orchestrator/gke/...
  • TestGeneratePathwaysManifest_MTCRejectsStaleCustomTemplate
  • TestValidateJobSetManifest_PodSpecVolumes, TestValidateJobSetManifest_DuplicateMountPathNormalized, TestValidatePathwaysMTCManifest_CleansMountPath

[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 enhances the reliability and configuration of Multi-Tier Checkpointing (MTC) for GKE jobs. It moves away from fragile string-based manifest checks toward robust structural validation of PodSpecs. Additionally, it streamlines MTC deployment by introducing a dedicated configuration resource in the Cluster Toolkit and improves safety by adding comprehensive pre-submission checks for flag consistency and mount conflicts.

Highlights

• Structural MTC Validation: Replaced substring-based validation of Pathways MTC manifests with structural PodSpec validation to prevent broken volume wiring.
• MTC Configuration Management: Introduced a dedicated checkpoint_configuration setting in the kubectl-apply module to manage the CheckpointConfiguration resource, replacing manual manifest application.
• Validation and Safety: Added strict validation for MTC flags, mount collisions, and compatibility checks between the MTC driver DaemonSet and target nodes.
• Infrastructure Migration: Simplified MTC IAM by removing legacy service account annotations and moving to direct Workload Identity Federation.

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 refactors and enhances the Multi-Tier Checkpointing (MTC) feature integration for GKE. Key changes include migrating from a legacy Workload Identity annotation model to a direct IAM principal binding, introducing a structured checkpoint_configuration setting in the kubectl-apply module to manage the CheckpointConfiguration CRD, and adding robust pre-flight validations (such as checking DaemonSet readiness, taint/toleration compatibility, and volume mount collisions with built-in Pathways paths). Additionally, the Pathways job template was updated to natively define MTC and colocated sidecar volumes on worker pods, and corresponding documentation, examples, and integration tests were updated. I have no feedback to provide as no review comments were submitted.

[shubpal07]

LGTM overall! Replacing the substring check with structured PodSpec validation is a much cleaner and safer approach.

Left some inline non-blocking suggestions.

[shubpal07]

LGTM