Termination message capture

[abhijeet-dhumal]

Add Termination Message-Based Final Metrics Capture

Summary

Add robust final metrics capture by reading pod termination messages, providing a reliable fallback when HTTP metrics endpoints become unavailable during pod shutdown. Includes comprehensive unit tests.

This is a 2nd part of #30, I'm targeting this entire fix into smaller atomic fixes via separate PRs here, : #32, #33, #34

Problem

Training progress metrics are currently captured via HTTP polling during execution. However, when pods terminate (normal completion, failure, preemption, OOM), the HTTP endpoint becomes unavailable before final metrics can be captured, leading to:

• Lost final progress updates (especially for epoch-based training)
• Inaccurate completion status
• Missing failure metrics when training crashes

Solution

Implement termination message-based capture as specified in the Kubeflow Training SDK:

1. Add CaptureMetricsFromTerminationMessage() function

• Update PollAndUpdateFinalProgress() to read from terminated pod containers
• Improve IsFinalStatusCaptured() with better grace period logic (3 polling cycles + buffer)
• Fix patch handling in PollAndUpdateProgress() for safer updates
• Support both "node" and "trainer" container names

2. Unit tests

Checklist:

• Docs included if any changes are user facing

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

✨ Finishing touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[robert-bell]

Thanks for your comment on #30. That makes sense.

Should we replace "node" with constants.Node here? And when would the container be called "trainer"?

[abhijeet-dhumal]

Thanks @robert-bell great catch! 🙌

[abhijeet-dhumal]

I have address this concern please check!

[robert-bell]

lgtm thanks Abhijeet!

[robert-bell]

lgtm thanks for this.

Just a minor comment which should be addressed, but I'm happy for you to merge this without another round of review.

[astefanutti]

/lgtm

Thanks!