[Neelabh94]

Summary

When --pathways-elastic-slices is 0 (non-elastic / rigid training), a single failed or preempted TPU worker pod cannot recover by restarting in-place while the rest of the slice continues running; the entire slice (and JobSet) must be torn down and recreated together so all workers re-handshake with the Pathways resource manager.

Previously, --pathways-max-slice-restarts unconditionally defaulted to 4000, rendering backoffLimitPerIndex: 4000 and backoffLimit: 4000 * nodesPerSlice on the worker Job even when --pathways-elastic-slices was 0. On a non-elastic job, a preempted worker pod was replaced in-place up to 4,000 times (DEADLINE_EXCEEDED: Worker failed to join a slice within 15m), preventing the worker Job from entering the Failed state and preventing the JobSet failurePolicy (--restarts) from triggering.

Changes

• cmd/job/submit.go:
  • Change the default of --pathways-max-slice-restarts to 0, and default pathways.MaxSliceRestarts to 4000 (defaultElasticMaxSliceRestarts) when pathways.ElasticSlices > 0 and --pathways-max-slice-restarts was not explicitly passed.
  • Validate that --pathways-elastic-slices and --pathways-max-slice-restarts are non-negative (>= 0).
  • Check hasChangedPathwaysFlags(cmd) so passing --pathways-max-slice-restarts 0 or --pathways-elastic-slices 0 without --pathways still returns pathways flags specified but --pathways flag is missing.
• cmd/job/submit_test.go:
  • Add table-driven unit tests covering default and explicit --pathways-max-slice-restarts values across --pathways-elastic-slices == 0 and > 0, negative flag values, and --pathways-* flags passed with 0 without --pathways.
• pkg/orchestrator/orchestrator.go, pkg/orchestrator/gke/manifest_generator_test.go, docs/gcluster_job_guide.md:
  • Update comments, test case names, and flag documentation to reflect the conditional default (0 when --pathways-elastic-slices is 0, or 4000 when --pathways-elastic-slices > 0).

Testing

• KUBECONFIG=/dev/null go test -race -count=1 ./cmd/job/... ./pkg/orchestrator/...
• pre-commit run --files cmd/job/submit.go cmd/job/submit_test.go pkg/orchestrator/orchestrator.go pkg/orchestrator/gke/manifest_generator_test.go docs/gcluster_job_guide.md

[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 modifies the default behavior of the --pathways-max-slice-restarts flag to better support non-elastic (rigid) training jobs. Previously, a default value of 4000 prevented non-elastic jobs from entering a failed state when a worker pod failed, as the system would repeatedly attempt to restart the pod in-place. The changes introduce a conditional default that sets the restart limit to 0 for non-elastic jobs, while maintaining the existing behavior for elastic jobs, alongside stricter flag validation and updated documentation.

Highlights

• Conditional Default for --pathways-max-slice-restarts: Updated the default value of --pathways-max-slice-restarts to 0 when --pathways-elastic-slices is 0, and to 4000 when elastic slices are enabled, ensuring non-elastic jobs can fail correctly.
• Flag Validation: Added validation to ensure --pathways-elastic-slices and --pathways-max-slice-restarts are non-negative and improved checks to prevent accidental use of pathways flags without the --pathways flag.
• Testing and Documentation: Expanded unit tests in cmd/job/submit_test.go to cover default scenarios and edge cases, and updated documentation in docs/gcluster_job_guide.md to reflect the new conditional default behavior.

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 updates the default behavior of the '--pathways-max-slice-restarts' flag to dynamically default to 0 when '--pathways-elastic-slices' is 0, and 4000 when '--pathways-elastic-slices' is greater than 0. It also adds validation to ensure these flags are non-negative and updates the associated tests and documentation. The feedback suggests a minor improvement in 'hasChangedPathwaysFlags' to use 'cmd.Flags().Visit' instead of 'VisitAll' for more idiomatic and efficient flag iteration.

[agrawalkhushi18]

Defaulting --pathways-max-slice-restarts to 0 for non-elastic jobs (--pathways-elastic-slices == 0) and 4000 for elastic jobs (--pathways-elastic-slices > 0) makes sense for Pod-level preemptions/evictions.

A few things worth addressing before merge:

1. Worker Pod restartPolicy: OnFailure in pathways_jobset.tmpl still restarts crashed worker containers in-place when MaxSliceRestarts == 0
   In pkg/orchestrator/gke/templates/pathways_jobset.tmpl (lines 222–242), the worker Job sets:

        backoffLimit: {{.WorkerBackoffLimit}}
        backoffLimitPerIndex: {{.WorkerBackoffLimitPerIndex}}
        podReplacementPolicy: Failed
        maxFailedIndexes: 0
        template:
          spec:
            restartPolicy: OnFailure

In Kubernetes Indexed Jobs, backoffLimitPerIndex only counts Pods that transition to Phase: Failed (e.g., when a TPU node is preempted/evicted). If a pathways-worker container crashes on a live node (e.g., SIGSEGV, container OOM, or DEADLINE_EXCEEDED: Worker failed to join a slice within 15m), restartPolicy: OnFailure causes kubelet to restart the pathways-worker container in-place inside the same Running Pod (CrashLoopBackOff) without the Pod ever entering Phase: Failed—so backoffLimitPerIndex: 0 and the JobSet failurePolicy never trigger.
When WorkerBackoffLimitPerIndex == 0 (or whenever backoffLimitPerIndex is used), should the worker Pod's restartPolicy be Never so container crashes immediately fail the Pod and trigger RestartJobSet?

{{- if eq .WorkerBackoffLimitPerIndex 0 }}
            restartPolicy: Never
{{- else }}
            restartPolicy: OnFailure
{{- end }}

2. Should --pathways-headless also default --pathways-max-slice-restarts to 4000?
   In pathwaysutils (experimental/gke/jobset.py), max_slice_restarts defaults high in headless mode as well so long-running headless Pathways clusters survive worker preemptions and client reconnections without tearing down the JobSet. With the current check in cmd/job/submit.go:

if pathways.ElasticSlices > 0 && !cmd.Flags().Changed("pathways-max-slice-restarts") {
    pathways.MaxSliceRestarts = defaultElasticMaxSliceRestarts
}

running gcluster job submit --pathways --pathways-headless (where --pathways-elastic-slices is 0) will now default --pathways-max-slice-restarts to 0. Should this check be if (pathways.ElasticSlices > 0 || pathways.Headless) && !cmd.Flags().Changed("pathways-max-slice-restarts")?

[Neelabh94]

Thanks for the thorough review, @agrawalkhushi18! Addressed both points in 9448434:

1. Worker Pod restartPolicy: Updated pkg/orchestrator/gke/templates/pathways_jobset.tmpl to render restartPolicy: Never when WorkerBackoffLimitPerIndex == 0 and restartPolicy: OnFailure otherwise, and added restartPolicy assertions in TestGeneratePathwaysManifest_MaxSliceRestarts.
2. --pathways-headless default: Updated runSubmitCmd in cmd/job/submit.go so --pathways-max-slice-restarts also defaults to 4000 when --pathways-headless is set (unless --pathways-max-slice-restarts is explicitly passed), along with unit tests in cmd/job/submit_test.go and updated flag documentation.

[agrawalkhushi18]

Thanks for addressing the comments, LGTM.