Make multi-host result reporting deterministic - #316
Merged
divyashreepathihalli merged 3 commits intoAug 19, 2026
Merged
Conversation
Every pod of a Pathways/LeaderWorkerSet job got the identical command,
including `--result-gcs gs://{bucket}/{job_id}/result.pkl`, and the
runner uploaded its result unconditionally. Two consequences: the value
returned to the client was whichever host wrote last, and a worker that
raised after the leader uploaded a success payload surfaced as a
`_false_success_error` carrying pod exit codes instead of the worker's
exception.
The runner now resolves its process index from TPU_WORKER_ID (falling
back to LWS_WORKER_INDEX, then to leader for single-host jobs) and:
- only the leader writes result.pkl, so the returned value is always
process 0's;
- a non-leader writes result-worker-<index>.pkl only when it fails, and
discards its return value without serializing it on success;
- both paths tag the payload with host_index.
`JobHandle.result()` aggregates those payloads for Pathways jobs: when
the leader claims success on a FAILED job, or wrote no payload at all,
the lowest-indexed failing host's exception is re-raised with its remote
traceback and a note naming the other failing hosts. Pod exit summaries
remain the fallback for hosts the kubelet killed before they could
report. Single-pod backends skip the listing entirely.
Adds unit tests for the runner gating and the client aggregation, plus
fake-gcs-server contract tests for the new storage helpers, and updates
docs/guides/distributed_training.md and AGENTS.md.
Rewrites the distributed-training section added with the previous commit in short active sentences, and removes the idioms and the contractions it carried. Splits the "read the pod logs" advice into two cases. Adds the matching entries to the troubleshooting page: - a new "A multi-host job fails, but the leader is successful" section that states which host reports the error and what the local error contains; - a note under "Job failed but no result payload was found" that the message now appears only when no host of the job wrote a record; - a note under "Job FAILS but logs look fine" that `kinetic jobs logs` returns the leader pod logs only, which can hide the failing host.
JyotinderSingh
requested review from
divyashreepathihalli and
jeffcarp
as code owners
August 15, 2026 19:07
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #316 +/- ##
=======================================
Coverage ? 36.73%
=======================================
Files ? 83
Lines ? 15532
Branches ? 0
=======================================
Hits ? 5706
Misses ? 9826
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request implements multi-host result ownership and failure aggregation for Pathways jobs. Under this new model, only the leader (process 0) writes the successful result payload, while non-leader hosts write failure payloads to worker-specific GCS blobs only if they fail. The client then aggregates these payloads and deterministically raises the exception of the lowest-indexed failing host. The feedback suggests optimizing the worker failure retrieval in kinetic/jobs.py by stopping the download loop as soon as the first valid failure payload is found, preventing sequential download bottlenecks on large multi-host jobs.
Review feedback on keras-team#316: the client downloaded every per-host failure payload but used only the first one's contents — the rest contributed their index to a note. One collective timeout leaves a payload on every host, so a large slice paid N sequential downloads to name N-1 hosts. `_worker_failure_error` now stops at the first payload that reports a failure and takes the remaining indices straight from the listing, which already carries them. A non-leader writes its blob only when it fails, so presence in the listing is the signal. The scan still skips a leading payload it cannot read, or one that claims success, so a corrupt blob on the lowest-indexed host does not hide the next host's error. `_worker_failures` collapses into `_list_worker_results`, which keeps the guarded listing. The note now reads "Other hosts that also reported a failure", matching what the listing alone proves.
divyashreepathihalli
approved these changes
Aug 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Every pod of a Pathways/LeaderWorkerSet job got the identical command, including
--result-gcs gs://{bucket}/{job_id}/result.pkl, and the runner uploaded its result unconditionally. Two consequences: the value returned to the client was whichever host wrote last, and a worker that raised after the leader uploaded a success payload surfaced as a_false_success_errorcarrying pod exit codes instead of the worker's exception.The runner now resolves its process index from TPU_WORKER_ID (falling back to LWS_WORKER_INDEX, then to leader for single-host jobs) and:
JobHandle.result()aggregates those payloads for Pathways jobs: when the leader claims success on a FAILED job, or wrote no payload at all, the lowest-indexed failing host's exception is re-raised with its remote traceback and a note naming the other failing hosts. Pod exit summaries remain the fallback for hosts the kubelet killed before they could report. Single-pod backends skip the listing entirely.Adds unit tests for the runner gating and the client aggregation, plus fake-gcs-server contract tests for the new storage helpers, and updates docs/guides/distributed_training.md and AGENTS.md.
Contributor Agreement
Please check all boxes below before submitting your PR for review: